from rich.console import Console
from dubbing.models import ProjectPaths, StepStatus, Language
from dubbing.steps import (
    ExtractAudioStep,
    ASRStep,
    TranslateStep,
    TTSStep,
    MixStep,
    FinalizeStep,
)

console = Console()


class Pipeline:
    def __init__(
        self,
        project_name: str,
        model_name: str = "gemini-2.0-flash-lite",
        language: Language = Language.RU,
    ):
        from dubbing.config import settings

        project_dir = settings.projects_dir / project_name
        self.paths = ProjectPaths.from_project(project_dir, language)
        self.model_name = model_name
        self.language = language
        self.steps = self._create_steps()

    def _create_steps(self) -> list:
        return [
            ExtractAudioStep(self.paths),
            ASRStep(self.paths),
            TranslateStep(self.paths, self.model_name, self.language),
            TTSStep(self.paths, self.language),
            MixStep(self.paths),
            FinalizeStep(self.paths),
        ]

    def get_cache_status(self) -> list[StepStatus]:
        return [
            StepStatus(
                name=step.name,
                cached=step.is_cached(),
                output_exists=step.is_cached(),
            )
            for step in self.steps
        ]

    def clean_from_step(self, step_index: int) -> None:
        for step in self.steps[step_index:]:
            step.clean()

    async def run(self, rebuild_from: int = -1) -> None:
        if rebuild_from >= 0:
            console.print(
                f"[yellow]Cleaning from step: {self.steps[rebuild_from].name}[/]"
            )
            self.clean_from_step(rebuild_from)

        for i, step in enumerate(self.steps):
            if step.is_cached():
                console.print(f"[dim]Skipping {step.name} (cached)[/]")
                continue

            console.print(
                f"\n[bold cyan]Step {i + 1}/{len(self.steps)}: {step.name}[/]"
            )
            await step.run()

        console.print("\n[bold green]Pipeline completed![/]")