from __future__ import annotations

from recoil.pipeline._lib import dispatch_payload as dp
from recoil.pipeline._lib.cinema_loader import render_cinema_tokens
from recoil.pipeline._lib.plan_loader import CanonicalShot, CharacterEntry
from recoil.pipeline._lib.prompt_engine import (
    _strip_focal_mm_tokens_for_model,
    bind_named_prose,
)
from recoil.pipeline._lib.shot_primitive import primitive_from_payload_context


MODE_ID = "seventies_new_wave_arriflex_cookes4"
MODEL_ID = "seeddance-2.0"
REF_MANIFEST = {"identity_1": 1, "identity_2": 2}


def _cinema_token_string(mode_id: str = MODE_ID) -> str:
    return _strip_focal_mm_tokens_for_model(
        render_cinema_tokens(mode_id, MODEL_ID).lower(),
        MODEL_ID,
    )


def _shot(
    shot_id: str,
    char_id: str,
    *,
    duration_s: float = 3.0,
    cinematography: dict | None = None,
    spoken_line: str | None = None,
) -> CanonicalShot:
    display_name = char_id.replace("_", " ").title()
    raw = {
        "shot_id": shot_id,
        "scene_index": 2,
        "duration_s": duration_s,
        "shot_type": "OTS",
        "camera_side": "B",
        "screen_direction": "left-to-right",
        "source_text": f"{display_name} moves through pressure.",
        "asset_data": {
            "location_id": "tartarus_med_bay",
            "characters": [{"char_id": char_id}],
        },
        "prompt_data": {
            "shot_type": "OTS",
            "action_line": f"{display_name} crosses the med bay.",
            "emotion_line": "Breath held tight.",
        },
        "routing_data": {
            "target_editorial_duration_s": duration_s,
            "has_dialogue": bool(spoken_line),
        },
        "refs": {
            "characters": {
                char_id: {
                    "display_name": display_name,
                    "aliases": [display_name],
                }
            }
        },
    }
    if cinematography is not None:
        raw["cinematography"] = cinematography
    if spoken_line is not None:
        raw["audio_data"] = {
            "dialogue": [{"speaker": char_id, "text": spoken_line}]
        }
    return CanonicalShot(
        shot_id=shot_id,
        scene_index=2,
        sequence_id=None,
        pipeline="video",
        previs_model="gemini-3-pro-image-preview",
        video_model=MODEL_ID,
        location_id="tartarus_med_bay",
        characters=[CharacterEntry(char_id=char_id)],
        shot_type="OTS",
        duration_s=duration_s,
        is_env_only=False,
        has_dialogue=bool(spoken_line),
        aspect_ratio="9:16",
        raw=raw,
        cinematography=cinematography,
    )


def _ctx(*, cinematography: dict | None = None, spoken_line: str | None = None):
    jade = _shot("EP002_SH01", "JADE", cinematography=cinematography)
    wren = _shot("EP002_SH02", "WREN", spoken_line=spoken_line)
    return dp.PayloadContext(
        project="tartarus",
        modality="r2v_multi",
        shot_id="EP002_PASS_009",
        shot=jade,
        batch_shots=[jade, wren],
        model_id=MODEL_ID,
        bible={},
    )


def _authored_prose() -> str:
    return (
        "[0:00-0:03] Jade pushes through Tartarus Med Bay as the camera tracks "
        "her tight breath.\n"
        "[0:03-0:06] Wren catches Jade by the sleeve as the lens racks onto "
        "his hand."
    )


def _patch_payload_io(monkeypatch, project_config: dict | None = None) -> None:
    dp._project_config_cache.clear()
    monkeypatch.setattr(dp, "load_project_config", lambda _project: project_config or {})
    monkeypatch.setattr(
        dp,
        "_collect_reference_images",
        lambda *args, **kwargs: (
            ["/tmp/jade.png", "/tmp/wren.png"],
            dict(REF_MANIFEST),
        ),
    )


def test_directed_prose_r2v_multi_appends_cinema_style_after_final_beat(
    monkeypatch,
) -> None:
    _patch_payload_io(monkeypatch)
    monkeypatch.setattr(dp, "author_pass", lambda *args, **kwargs: _authored_prose())

    payload = dp.build_unified_payload(_ctx(cinematography={"mode": MODE_ID}))
    prompt = payload["prompt"]

    assert "Style:" in prompt
    assert "kodak 5247 color negative film stock emulation" in prompt
    assert "warm desaturated 70s print look" in prompt
    assert _cinema_token_string() in prompt
    assert prompt.index("Style:") > prompt.rindex("[0:03-0:06]")


def test_directed_prose_without_cinema_or_film_stock_keeps_bound_prose_exact(
    monkeypatch,
) -> None:
    _patch_payload_io(monkeypatch)
    authored = _authored_prose()
    ctx = _ctx()
    monkeypatch.setattr(dp, "author_pass", lambda *args, **kwargs: authored)
    monkeypatch.setattr(dp, "_get_seeddance_film_stock", lambda _project_config: "")

    payload = dp.build_unified_payload(ctx)
    primitive = primitive_from_payload_context(
        ctx,
        ref_manifest=REF_MANIFEST,
        segment_timestamps=[(0.0, 3.0), (3.0, 6.0)],
    )
    expected = bind_named_prose(
        authored,
        primitive,
        REF_MANIFEST,
        modality="r2v_multi",
    ).text

    assert payload["prompt"] == expected
    assert "Style:" not in payload["prompt"]
    assert not payload["prompt"].rstrip().endswith("Style:")


def test_deterministic_template_path_does_not_double_append_cinema_tokens(
    monkeypatch,
) -> None:
    _patch_payload_io(monkeypatch, project_config={"cinema_mode": MODE_ID})
    token_string = _cinema_token_string()
    template_prompt = f"Template prose.\n\nStyle: {token_string}. Cinematic, photorealistic."
    monkeypatch.setenv("PROSE_AUTHOR_FALLBACK", "1")
    monkeypatch.setattr(
        dp,
        "get_builder",
        lambda *args, **kwargs: lambda *a, **k: template_prompt,
    )

    payload = dp.build_unified_payload(_ctx(cinematography={"mode": MODE_ID}))

    assert payload["prompt"].count(token_string) == 1


def test_dialogue_formatting_unit_survives_directed_prose_binding(
    monkeypatch,
) -> None:
    _patch_payload_io(monkeypatch)
    line = "You are currently listed as Expendable."
    authored = (
        "[0:00-0:03] Jade pushes through Tartarus Med Bay as the camera tracks "
        "her tight breath.\n"
        '[0:03-0:06] Wren, flat and unweighted, says: "'
        f'{line}" as the lens holds on his face.'
    )
    monkeypatch.setattr(dp, "author_pass", lambda *args, **kwargs: authored)

    payload = dp.build_unified_payload(_ctx(spoken_line=line))
    prompt = payload["prompt"]
    unit = f'@Image2, flat and unweighted, says: "{line}"'

    assert unit in prompt
    assert prompt.index("flat and unweighted") < prompt.index(f'"{line}"')
    assert line in prompt