from pathlib import Path
from typing import Dict, Iterable, List, Optional

from backend.types import RenderArtifact, RenderJob, VoiceConfig


class RenderPipeline:
    def __init__(self, session_root: Path, synthesizer=None) -> None:
        self.session_root = Path(session_root)
        self.session_root.mkdir(parents=True, exist_ok=True)
        self.synthesizer = synthesizer
        self._jobs: Dict[str, RenderJob] = {}
        self._cancelled: set[str] = set()
        self._paused: set[str] = set()

    def get_job(self, session_id: str) -> RenderJob:
        return self._jobs.setdefault(session_id, RenderJob(session_id=session_id))

    def cancel(self, session_id: str) -> None:
        self._cancelled.add(session_id)
        self.get_job(session_id).status = "cancelled"

    def pause(self, session_id: str) -> Dict[str, object]:
        self._paused.add(session_id)
        self.get_job(session_id).status = "paused"
        return {"type": "paused", "session_id": session_id}

    def resume(self, session_id: str) -> Dict[str, object]:
        self._paused.discard(session_id)
        self.get_job(session_id).status = "running"
        return {"type": "resumed", "session_id": session_id}

    def render(
        self,
        *,
        session_id: str,
        book: Dict[str, object],
        chapters: List[Dict[str, object]],
        voice_config: Dict[str, object],
        diffusion_steps: int,
        speed: float,
        synthesizer=None,
    ) -> Iterable[Dict[str, object]]:
        selected = [chapter for chapter in chapters if chapter.get("included", True)]
        job = self.get_job(session_id)
        job.status = "running"
        job.outputs.clear()
        synth = synthesizer or self.synthesizer
        if synth is None:
            raise ValueError("No synthesizer provided for render pipeline")

        yield {
            "type": "started",
            "session_id": session_id,
            "total_chapters": len(selected),
            "book_title": book.get("title"),
            "backend": "local",
            "model": VoiceConfig.from_dict(voice_config).model,
        }

        voice = VoiceConfig.from_dict(voice_config)
        render_dir = self.session_root / session_id / "renders"
        render_dir.mkdir(parents=True, exist_ok=True)

        for index, chapter in enumerate(selected):
            if session_id in self._cancelled:
                yield {"type": "cancelled", "session_id": session_id, "backend": "local", "model": voice.model}
                self._cancelled.discard(session_id)
                job.status = "cancelled"
                return

            chapter_id = str(chapter["id"])
            job.current_chapter_id = chapter_id
            yield {
                "type": "chapter_started",
                "session_id": session_id,
                "chapter_id": chapter_id,
                "chapter_title": chapter["title"],
                "chapter_index": index,
                "overall_progress": index / max(1, len(selected)),
                "backend": "local",
                "model": voice.model,
            }

            output_path = render_dir / f"{index + 1:03d}-{chapter_id}.wav"
            result = synth.synthesize(
                text=str(chapter["text"]),
                output_path=output_path,
                voice_config=voice,
                diffusion_steps=diffusion_steps,
                speed=speed,
            )

            yield {
                "type": "chapter_progress",
                "session_id": session_id,
                "chapter_id": chapter_id,
                "percent": 100,
                "backend": result.get("backend", "local"),
                "model": result.get("model", voice.model),
            }

            artifact = RenderArtifact(
                path=output_path,
                duration_seconds=int(result.get("duration_seconds", 0)),
                chapter_id=chapter_id,
            )
            job.outputs.append(artifact)
            yield {
                "type": "chapter_done",
                "session_id": session_id,
                "chapter_id": chapter_id,
                "duration_seconds": artifact.duration_seconds,
                "overall_progress": (index + 1) / max(1, len(selected)),
                "output_path": str(output_path),
                "backend": result.get("backend", "local"),
                "model": result.get("model", voice.model),
            }

        job.status = "completed"
        yield {
            "type": "completed",
            "session_id": session_id,
            "outputs": [str(artifact.path) for artifact in job.outputs],
            "backend": "local",
            "model": voice.model,
        }