from enum import Enum
from pydantic import BaseModel
from pathlib import Path


class Language(str, Enum):
    RU = "ru"
    EN = "en"
    EN_RU = "en_ru"  # Chinese -> English -> Russian


class Segment(BaseModel):
    start: int
    end: int
    text: str


class TranslatedSegment(Segment):
    translated: str


class Translation(BaseModel):
    id: int
    translated: str


class TranslationBatch(BaseModel):
    translations: list[Translation]


class TTSSegment(TranslatedSegment):
    audio_path: Path
    audio_duration_ms: float


class ProjectPaths(BaseModel):
    root: Path
    source_video: Path
    source_audio: Path
    segments_json: Path
    translated_json: Path
    tts_dir: Path
    dubbed_audio: Path
    result_video: Path

    model_config = {"arbitrary_types_allowed": True}

    @classmethod
    def from_project(
        cls, project_dir: Path, language: Language = Language.RU
    ) -> "ProjectPaths":
        lang = language.value
        return cls(
            root=project_dir,
            source_video=project_dir / "source.mp4",
            source_audio=project_dir / "source.mp3",
            segments_json=project_dir / "segments.json",
            translated_json=project_dir / f"translated_{lang}.json",
            tts_dir=project_dir / f"tts_{lang}",
            dubbed_audio=project_dir / f"dubbed_{lang}.mp3",
            result_video=project_dir / f"result_{lang}.mkv",
        )


class StepStatus(BaseModel):
    name: str
    cached: bool
    output_exists: bool