"""Tests for pipeline runner — checkpointing and resume."""

import json
from pathlib import Path
from typing import Any
from unittest.mock import MagicMock, patch

import pytest

from tests.conftest import prefill_download
from yt_dbl.config import Settings
from yt_dbl.pipeline.base import StepValidationError
from yt_dbl.pipeline.runner import PipelineRunner, load_state, save_state
from yt_dbl.schemas import STEP_DIRS, PipelineState, Segment, StepName, StepStatus


def _fake_separation_factory(sep_dir: Path) -> Any:
    """Return a side_effect callable that creates fake separation outputs."""

    def _fake(audio_path: Path) -> None:
        (sep_dir / "vocals.wav").write_bytes(b"fake-vocals")
        (sep_dir / "background.wav").write_bytes(b"fake-background")

    return _fake


def _fake_transcription_factory() -> dict[str, Any]:
    """Return mock side_effects for _run_asr and _run_alignment."""
    from yt_dbl.schemas import Segment, Word

    fake_asr_segments = [
        {"start": 0.0, "end": 3.0, "speaker_id": 0, "text": "Hello world."},
        {"start": 4.0, "end": 6.0, "speaker_id": 1, "text": "Goodbye."},
    ]

    fake_segments = [
        Segment(
            id=0,
            text="Hello world.",
            start=0.0,
            end=3.0,
            speaker="SPEAKER_00",
            language="en",
            words=[Word(text="Hello", start=0.0, end=1.5), Word(text="world.", start=1.5, end=3.0)],
        ),
        Segment(
            id=1,
            text="Goodbye.",
            start=4.0,
            end=6.0,
            speaker="SPEAKER_01",
            language="en",
            words=[Word(text="Goodbye.", start=4.0, end=6.0)],
        ),
    ]

    def _fake_asr(vocals_path: Path) -> list[dict[str, Any]]:
        return fake_asr_segments

    def _fake_alignment(
        vocals_path: Path,
        raw_segments: list[dict[str, Any]],
        detected_lang: str = "en",
    ) -> list[Segment]:
        return fake_segments

    return {
        "asr": _fake_asr,
        "alignment": _fake_alignment,
    }


def _fake_translation(
    segments: list[Segment],
    target_language: str,
    source_language: str = "auto-detected",
) -> dict[int, str]:
    """Fake translate that returns deterministic translations."""
    return {seg.id: f"[{target_language}] {seg.text}" for seg in segments}


def _fake_postprocess_segment(input_path: Path, output_path: Path, **_kwargs: Any) -> Path:
    """Mock postprocess_segment that creates the output file."""
    output_path.parent.mkdir(parents=True, exist_ok=True)
    output_path.write_bytes(b"fake-audio")
    return output_path


def _fake_tts_model() -> Any:
    """Return a fake TTS model whose generate() returns dummy audio."""
    import numpy as np

    model = MagicMock()
    result = MagicMock()
    result.audio = np.zeros(12000, dtype=np.float32)
    result.sample_rate = 12000
    model.generate.return_value = [result]
    return model


def _pipeline_patches(sep_dir: Path) -> Any:
    """Context manager stack for all pipeline mocks (separate+transcribe+translate+synthesize)."""
    from contextlib import ExitStack

    fakes = _fake_transcription_factory()

    stack = ExitStack()
    stack.enter_context(
        patch(
            "yt_dbl.pipeline.separate.SeparateStep._run_separation",
            side_effect=_fake_separation_factory(sep_dir),
        )
    )
    stack.enter_context(
        patch(
            "yt_dbl.pipeline.transcribe.TranscribeStep._run_asr",
            side_effect=fakes["asr"],
        )
    )
    stack.enter_context(
        patch(
            "yt_dbl.pipeline.transcribe.TranscribeStep._run_alignment",
            side_effect=fakes["alignment"],
        )
    )
    stack.enter_context(
        patch(
            "yt_dbl.pipeline.translate.TranslateStep._translate_all",
            side_effect=_fake_translation,
        )
    )
    # Synthesize mocks: skip real TTS + ffmpeg
    stack.enter_context(
        patch(
            "yt_dbl.pipeline.synthesize.SynthesizeStep._load_tts_model",
            return_value=_fake_tts_model(),
        )
    )
    stack.enter_context(
        patch(
            "yt_dbl.pipeline.synthesize.SynthesizeStep._run_tts",
            return_value=__import__("numpy").zeros(12000, dtype=__import__("numpy").float32),
        )
    )
    stack.enter_context(
        patch("yt_dbl.utils.audio_processing.run_ffmpeg"),
    )
    stack.enter_context(
        patch(
            "yt_dbl.pipeline.synthesize.postprocess_segment",
            side_effect=_fake_postprocess_segment,
        ),
    )
    stack.enter_context(
        patch(
            "yt_dbl.pipeline.synthesize.get_audio_duration",
            return_value=1.0,
        ),
    )
    stack.enter_context(
        patch("yt_dbl.pipeline.synthesize.SynthesizeStep._save_wav"),
    )
    # Assemble mocks: skip real ffmpeg for final muxing
    stack.enter_context(
        patch("yt_dbl.pipeline.assemble.run_ffmpeg"),
    )
    return stack


class TestCheckpoints:
    def test_save_and_load(self, tmp_path: Path) -> None:
        cfg = Settings(work_dir=tmp_path / "work")
        state = PipelineState(video_id="test123", url="https://example.com", target_language="ru")
        state.get_step(StepName.DOWNLOAD).status = StepStatus.COMPLETED

        save_state(state, cfg)
        loaded = load_state(cfg, "test123")

        assert loaded is not None
        assert loaded.video_id == "test123"
        assert loaded.get_step(StepName.DOWNLOAD).status == StepStatus.COMPLETED

    def test_load_nonexistent(self, tmp_path: Path) -> None:
        cfg = Settings(work_dir=tmp_path / "work")
        assert load_state(cfg, "nonexistent") is None

    def test_state_file_is_valid_json(self, tmp_path: Path) -> None:
        cfg = Settings(work_dir=tmp_path / "work")
        state = PipelineState(video_id="test123")
        path = save_state(state, cfg)

        data = json.loads(path.read_text())
        assert data["video_id"] == "test123"

    def test_save_no_leftover_tmp_files(self, tmp_path: Path) -> None:
        """After save_state completes, no .tmp files remain."""
        cfg = Settings(work_dir=tmp_path / "work")
        state = PipelineState(video_id="test123")
        save_state(state, cfg)

        job_dir = cfg.work_dir / "test123"
        tmp_files = list(job_dir.glob("*.tmp"))
        assert tmp_files == []

    def test_save_overwrites_atomically(self, tmp_path: Path) -> None:
        """A second save_state replaces the first without corruption."""
        cfg = Settings(work_dir=tmp_path / "work")
        state = PipelineState(video_id="test123", target_language="en")
        save_state(state, cfg)

        state.target_language = "fr"
        save_state(state, cfg)

        loaded = load_state(cfg, "test123")
        assert loaded is not None
        assert loaded.target_language == "fr"

    def test_save_cleans_tmp_on_error(self, tmp_path: Path) -> None:
        """If serialisation fails, no .tmp file is left behind."""
        cfg = Settings(work_dir=tmp_path / "work")
        state = PipelineState(video_id="test123")
        # Ensure job dir exists
        cfg.job_dir("test123")

        with (
            patch.object(
                PipelineState,
                "model_dump_json",
                side_effect=RuntimeError("boom"),
            ),
            pytest.raises(RuntimeError, match="boom"),
        ):
            save_state(state, cfg)

        job_dir = cfg.work_dir / "test123"
        tmp_files = list(job_dir.glob("*.tmp"))
        assert tmp_files == []


class TestPipelineRunner:
    def test_pipeline_from_separate(self, tmp_path: Path) -> None:
        """Run pipeline skipping download (which needs real yt-dlp)."""
        cfg = Settings(work_dir=tmp_path / "work", anthropic_api_key="sk-test")
        state = PipelineState(
            video_id="test123",
            url="https://youtube.com/watch?v=test123",
            target_language="ru",
        )
        state = prefill_download(state, cfg)
        save_state(state, cfg)

        sep_dir = cfg.step_dir("test123", STEP_DIRS[StepName.SEPARATE])

        runner = PipelineRunner(cfg)
        with _pipeline_patches(sep_dir):
            state = runner.run(state)

        assert state.next_step is None
        for step_name in [
            StepName.SEPARATE,
            StepName.TRANSCRIBE,
            StepName.TRANSLATE,
            StepName.SYNTHESIZE,
            StepName.ASSEMBLE,
        ]:
            assert state.get_step(step_name).status == StepStatus.COMPLETED

    def test_resume_skips_completed(self, tmp_path: Path) -> None:
        cfg = Settings(work_dir=tmp_path / "work", anthropic_api_key="sk-test")
        state = PipelineState(video_id="test123", url="https://example.com")
        state = prefill_download(state, cfg)

        sep_dir = cfg.step_dir("test123", STEP_DIRS[StepName.SEPARATE])

        runner = PipelineRunner(cfg)
        with _pipeline_patches(sep_dir):
            state = runner.run(state)

        # Mark last 3 as pending again
        for name in [StepName.TRANSLATE, StepName.SYNTHESIZE, StepName.ASSEMBLE]:
            state.get_step(name).status = StepStatus.PENDING
        save_state(state, cfg)

        loaded = load_state(cfg, "test123")
        assert loaded is not None
        assert loaded.next_step == StepName.TRANSLATE

    def test_from_step(self, tmp_path: Path) -> None:
        cfg = Settings(work_dir=tmp_path / "work", anthropic_api_key="sk-test")
        state = PipelineState(video_id="test123", url="https://example.com")
        state = prefill_download(state, cfg)

        sep_dir = cfg.step_dir("test123", STEP_DIRS[StepName.SEPARATE])

        runner = PipelineRunner(cfg)
        with _pipeline_patches(sep_dir):
            state = runner.run(state)
            state = runner.run(state, from_step=StepName.TRANSLATE)

        assert state.get_step(StepName.TRANSLATE).status == StepStatus.COMPLETED

    def test_step_failure_stops_pipeline(self, tmp_path: Path) -> None:
        """When a step raises, its status becomes FAILED and pipeline stops."""
        cfg = Settings(work_dir=tmp_path / "work", anthropic_api_key="sk-test")
        state = PipelineState(video_id="test123", url="https://example.com")
        state = prefill_download(state, cfg)

        runner = PipelineRunner(cfg)

        # Patch SeparateStep.run to raise
        with patch(
            "yt_dbl.pipeline.separate.SeparateStep.run",
            side_effect=RuntimeError("boom"),
        ):
            state = runner.run(state)

        assert state.get_step(StepName.SEPARATE).status == StepStatus.FAILED
        assert "boom" in state.get_step(StepName.SEPARATE).error
        # Steps after failure should still be pending
        assert state.get_step(StepName.TRANSCRIBE).status == StepStatus.PENDING

    def test_checkpoint_saved_on_success(self, tmp_path: Path) -> None:
        """After each step completes, state is persisted to disk."""
        cfg = Settings(work_dir=tmp_path / "work", anthropic_api_key="sk-test")
        state = PipelineState(video_id="test123", url="https://example.com")
        state = prefill_download(state, cfg)

        sep_dir = cfg.step_dir("test123", STEP_DIRS[StepName.SEPARATE])

        runner = PipelineRunner(cfg)
        with _pipeline_patches(sep_dir):
            state = runner.run(state)

        loaded = load_state(cfg, "test123")
        assert loaded is not None
        assert loaded.get_step(StepName.ASSEMBLE).status == StepStatus.COMPLETED

    def test_checkpoint_saved_on_failure(self, tmp_path: Path) -> None:
        """Failed state should also be persisted."""
        cfg = Settings(work_dir=tmp_path / "work", anthropic_api_key="sk-test")
        state = PipelineState(video_id="test123", url="https://example.com")
        state = prefill_download(state, cfg)

        runner = PipelineRunner(cfg)

        with patch(
            "yt_dbl.pipeline.separate.SeparateStep.run",
            side_effect=RuntimeError("disk full"),
        ):
            state = runner.run(state)

        loaded = load_state(cfg, "test123")
        assert loaded is not None
        assert loaded.get_step(StepName.SEPARATE).status == StepStatus.FAILED
        assert "disk full" in loaded.get_step(StepName.SEPARATE).error

    def test_step_timing_recorded(self, tmp_path: Path) -> None:
        """Completed steps record timing metadata."""
        cfg = Settings(work_dir=tmp_path / "work", anthropic_api_key="sk-test")
        state = PipelineState(video_id="test123", url="https://example.com")
        state = prefill_download(state, cfg)

        sep_dir = cfg.step_dir("test123", STEP_DIRS[StepName.SEPARATE])

        runner = PipelineRunner(cfg)
        with _pipeline_patches(sep_dir):
            state = runner.run(state)

        result = state.get_step(StepName.SEPARATE)
        assert result.started_at != ""
        assert result.finished_at != ""
        assert result.duration_sec >= 0.0


class TestEarlyValidation:
    def test_missing_api_key_raises(self, tmp_path: Path) -> None:
        """Pipeline refuses to start when API key is missing and translate hasn't run."""
        cfg = Settings(work_dir=tmp_path / "work", anthropic_api_key="")
        state = PipelineState(video_id="test123", url="https://example.com")
        state = prefill_download(state, cfg)

        runner = PipelineRunner(cfg)
        with pytest.raises(StepValidationError, match="Anthropic API key"):
            runner.run(state)

    def test_api_key_not_needed_when_translate_done(self, tmp_path: Path) -> None:
        """If translate is already COMPLETED, no API key is required."""
        cfg = Settings(work_dir=tmp_path / "work", anthropic_api_key="")
        state = PipelineState(video_id="test123", url="https://example.com")
        state = prefill_download(state, cfg)
        state.get_step(StepName.TRANSLATE).status = StepStatus.COMPLETED

        sep_dir = cfg.step_dir("test123", STEP_DIRS[StepName.SEPARATE])
        runner = PipelineRunner(cfg)
        with _pipeline_patches(sep_dir):
            state = runner.run(state)
        # Pipeline should run without raising


class TestModelCleanup:
    def test_models_unloaded_on_keyboard_interrupt(self, tmp_path: Path) -> None:
        """Models are freed even if a step raises KeyboardInterrupt."""
        cfg = Settings(work_dir=tmp_path / "work", anthropic_api_key="sk-test")
        state = PipelineState(video_id="test123", url="https://example.com")
        state = prefill_download(state, cfg)

        runner = PipelineRunner(cfg)

        with (
            patch.object(runner.model_manager, "unload_all") as mock_unload,
            patch(
                "yt_dbl.pipeline.separate.SeparateStep.run",
                side_effect=KeyboardInterrupt,
            ),
            pytest.raises(KeyboardInterrupt),
        ):
            runner.run(state)

        mock_unload.assert_called_once()
