from __future__ import annotations

import re

import pytest

from recoil.pipeline._lib.prompt_engine import (
    BindAssertionError,
    bind_named_prose,
)
from recoil.pipeline._lib.shot_primitive import ShotPrimitive


def _r2v_primitive() -> ShotPrimitive:
    return ShotPrimitive(
        shot_id="EP002_PASS_009",
        scene_index=2,
        shot_type="OTS",
        target_editorial_duration_s=6,
        intent="Jade and Wren cross the med bay under pressure.",
        camera_side="B",
        screen_direction="left-to-right",
        char_ids=["JADE", "WREN"],
        location_id="tartarus_med_bay",
        timing_segments=[
            {"start_s": 0.0, "end_s": 3.0, "intent": "Jade moves first."},
            {"start_s": 3.0, "end_s": 6.0, "intent": "Wren answers."},
        ],
        refs={
            "characters": {
                "JADE": {
                    "display_name": "Jade",
                    "aliases": ["Jade Cooper"],
                    "role": "protagonist",
                },
                "WREN": {
                    "display_name": "Wren",
                    "aliases": ["Wren Vale"],
                    "role": "support",
                },
            },
            "locations": {
                "tartarus_med_bay": {"display_name": "Tartarus Med Bay"},
            },
        },
    )


def test_bind_named_prose_r2v_replaces_character_names_and_preserves_locations() -> None:
    primitive = _r2v_primitive()
    authored = (
        "[0:00-0:03] Jade Cooper leans into Wren as the lens pushes through "
        "Tartarus Med Bay.\n"
        "[0:03-0:06] Wren Vale steadies Jade with a tight breath."
    )

    bound = bind_named_prose(
        authored,
        primitive,
        {"identity_1": 1, "identity_2": 2, "scene_1": 3},
        modality="r2v_multi",
    )

    assert "Jade" not in bound.text
    assert "Jade Cooper" not in bound.text
    assert "Wren" not in bound.text
    assert "Wren Vale" not in bound.text
    assert "Tartarus Med Bay" in bound.text
    assert "@Image1" in bound.text
    assert "@Image2" in bound.text
    assert bound.manifest == {"identity_1": 1, "identity_2": 2, "scene_1": 3}
    assert {int(n) for n in re.findall(r"@Image(\d+)", bound.text)} <= {
        int(v) for v in bound.manifest.values()
    }


def test_bind_named_prose_i2v_uses_live_payload_ref_keys_and_no_image_tokens() -> None:
    primitive = ShotPrimitive(
        shot_id="EP001_SH11",
        scene_index=3,
        shot_type="OTS",
        target_editorial_duration_s=5,
        intent="Jade crosses from the start frame to the end frame.",
        char_ids=["JADE"],
        refs={
            "start_frame": "/tmp/start.png",
            "end_frame": "/tmp/end.png",
            "characters": {
                "JADE": {"display_name": "Jade", "aliases": ["Jade Cooper"]},
            },
        },
    )

    bound = bind_named_prose(
        "Jade Cooper pulls herself upright as the corridor light changes.",
        primitive,
        {"identity_1": 1},
        modality="video_i2v",
    )

    assert bound.payload_refs == {
        "start_frame": "/tmp/start.png",
        "image_tail": "/tmp/end.png",
    }
    assert "@Image" not in bound.text
    assert "Jade" not in bound.text
    assert "Jade Cooper" not in bound.text


def test_bind_named_prose_i2v_rejects_leaked_image_tokens() -> None:
    primitive = ShotPrimitive(
        shot_id="EP001_SH11",
        scene_index=3,
        shot_type="OTS",
        target_editorial_duration_s=5,
        intent="Jade crosses from the start frame to the end frame.",
        char_ids=["JADE"],
        refs={"start_frame": "/tmp/start.png", "end_frame": "/tmp/end.png"},
    )

    with pytest.raises(BindAssertionError):
        bind_named_prose(
            "The subject moves from @Image1 into the final frame.",
            primitive,
            {"identity_1": 1},
            modality="video_i2v",
        )