from __future__ import annotations

import tempfile
import unittest
from pathlib import Path

from production_build_agents.run.segment_dependencies import (
    SegmentDependencyIntegrityError,
    load_verified_segment_dependencies,
)
from production_build_agents.run.segment_inputs import (
    load_verified_segment_inputs,
)
from production_build_agents.agents.segment.executor import _user_payload
from production_build_agents.agents.segment.provenance import (
    normalize_candidate_lineage,
)
from production_build_agents.contracts.models import ToolCallRecord
from production_build_agents.contracts.production_evaluation import (
    evaluate_segment_delivery,
)
from production_build_agents.contracts.document_identity import (
    exact_document_sha256,
)
from production_build_agents.contracts.production_execution_models import (
    ImmutableJsonRef,
    SegmentDependencyRef,
    SegmentPackage,
)
from production_build_agents.run.artifacts import (
    content_identity_for_path,
    file_sha256,
    seal_shared_visual_anchor,
)
from tests.support.production_fixtures import (
    output_profile,
    segment_bundle,
)
from tests.support.production_planner_fixtures import (
    build_production_planner_scenario,
)


def _ref(path: Path) -> ImmutableJsonRef:
    return ImmutableJsonRef(
        uri=str(path.resolve()),
        document_sha256=exact_document_sha256(path.read_bytes()),
    )


class SegmentDependencyTest(unittest.TestCase):
    def _records(self, root: Path):
        (
            source_package,
            _source_candidate,
            source_delivery,
            _validator_candidate,
            source_report,
            _accepted,
        ) = segment_bundle(1)
        package_path = root / "upstream-package.json"
        package_path.write_text(
            source_package.model_dump_json(),
            encoding="utf-8",
        )
        sealed_artifacts = []
        for artifact in source_delivery.artifacts:
            path = root / f"{artifact.artifact_id}.bin"
            path.write_bytes(artifact.artifact_id.encode())
            sealed_artifacts.append(
                artifact.model_copy(
                    update={
                        "uri": str(path.resolve()),
                        "source_uri": None,
                        **content_identity_for_path(path),
                    }
                )
            )
        source_delivery = source_delivery.model_copy(
            update={
                "segment_package_uri": str(package_path.resolve()),
                "segment_package_sha256": file_sha256(package_path),
                "artifacts": sealed_artifacts,
            }
        )
        delivery_path = root / "upstream-delivery.json"
        report_path = root / "upstream-report.json"
        delivery_path.write_text(
            source_delivery.model_dump_json(),
            encoding="utf-8",
        )
        report_path.write_text(
            source_report.model_dump_json(),
            encoding="utf-8",
        )
        dependency = SegmentDependencyRef(
            segment_id="Segment1",
            segment_delivery_ref=_ref(delivery_path),
            validation_report_ref=_ref(report_path),
        )

        (
            downstream_base,
            downstream_candidate,
            downstream_delivery,
            _validator_candidate,
            _report,
            _accepted,
        ) = segment_bundle(2)
        downstream = downstream_base.model_copy(
            update={"dependencies": [dependency]}
        )
        return (
            downstream,
            downstream_candidate,
            downstream_delivery,
            source_delivery,
            report_path,
        )

    def test_exact_refs_pass_report_and_full_agent_payload(self) -> None:
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            package, _candidate, _delivery, source, _report_path = (
                self._records(root)
            )
            verified = load_verified_segment_dependencies(package)
            self.assertEqual(len(verified), 1)
            sealed_package_artifacts = []
            for artifact in package.artifacts:
                path = root / f"selected-{artifact.artifact_id}.bin"
                path.write_bytes(artifact.artifact_id.encode())
                sealed_package_artifacts.append(
                    artifact.model_copy(
                        update={
                            "uri": str(path.resolve()),
                            "source_uri": None,
                            **content_identity_for_path(path),
                        }
                    )
                )
            global_delivery = build_production_planner_scenario(
                run_id="dependency-global"
            ).global_data_delivery.model_copy(
                update={"active_artifacts": sealed_package_artifacts}
            )
            global_delivery_path = root / "global-delivery.json"
            global_delivery_path.write_text(
                global_delivery.model_dump_json(),
                encoding="utf-8",
            )
            anchor_path = root / "shared-anchor.png"
            anchor_path.write_bytes(b"shared-anchor")
            package = package.model_copy(
                update={
                    "artifacts": sealed_package_artifacts,
                    "global_data_delivery_ref": _ref(global_delivery_path),
                    "shared_visual_anchor": seal_shared_visual_anchor(
                        anchor_path
                    ),
                }
            )
            verified_inputs = load_verified_segment_inputs(package)
            payload = _user_payload(
                package,
                executor_run_id="run:Segment2:v1:executor",
                verified_inputs=verified_inputs,
            )
            dependency_payload = payload[
                "verified_segment_inputs"
            ]["accepted_segments"][0]
            self.assertEqual(
                [
                    item["artifact_id"]
                    for item in dependency_payload["artifacts"]
                ],
                [item.artifact_id for item in source.artifacts],
            )
            self.assertTrue(
                all(
                    item["uri"] == source.artifacts[index].uri
                    for index, item in enumerate(
                        dependency_payload["artifacts"]
                    )
                )
            )
            self.assertTrue(
                all(
                    "source_uri" not in item
                    for item in dependency_payload["artifacts"]
                )
            )
            self.assertNotIn("tool_calls", dependency_payload)

    def test_tampered_or_failed_report_is_rejected_before_execution(
        self,
    ) -> None:
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            package, _candidate, _delivery, _source, report_path = (
                self._records(root)
            )
            report_path.write_text("{}", encoding="utf-8")
            with self.assertRaisesRegex(
                SegmentDependencyIntegrityError,
                "精确字节已变化",
            ):
                load_verified_segment_dependencies(package)

    def test_accepted_artifact_uri_cannot_conflict_with_global_input(
        self,
    ) -> None:
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            package, _candidate, _delivery, source, _report_path = (
                self._records(root)
            )
            package = package.model_copy(
                update={
                    "artifacts": [
                        package.artifacts[0].model_copy(
                            update={"uri": source.artifacts[0].uri}
                        )
                    ]
                }
            )
            with self.assertRaisesRegex(
                SegmentDependencyIntegrityError,
                "ID 和 URI 必须全局无歧义",
            ):
                load_verified_segment_dependencies(package)

    def test_failed_dependency_report_is_rejected_before_execution(
        self,
    ) -> None:
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            package, _candidate, _delivery, _source, report_path = (
                self._records(root)
            )
            report = package.dependencies[0].validation_report_ref
            from production_build_agents.contracts.production_models import (
                SegmentValidationReport,
            )

            failed = SegmentValidationReport.model_validate_json(
                report_path.read_bytes()
            ).model_copy(
                update={
                    "criterion_results": [
                        SegmentValidationReport.model_validate_json(
                            report_path.read_bytes()
                        ).criterion_results[0].model_copy(
                            update={"verdict": "FAIL"}
                        ),
                        *SegmentValidationReport.model_validate_json(
                            report_path.read_bytes()
                        ).criterion_results[1:],
                    ],
                    "verdict": "FAIL",
                }
            )
            report_path.write_text(
                failed.model_dump_json(),
                encoding="utf-8",
            )
            package = package.model_copy(
                update={
                    "dependencies": [
                        package.dependencies[0].model_copy(
                            update={
                                "validation_report_ref": _ref(report_path)
                            }
                        )
                    ]
                }
            )
            with self.assertRaisesRegex(
                SegmentDependencyIntegrityError,
                "不是已验收 PASS",
            ):
                load_verified_segment_dependencies(package)

    def test_dependency_media_is_an_allowed_formal_lineage_input(
        self,
    ) -> None:
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            package, candidate, delivery, source, _report_path = (
                self._records(root)
            )
            upstream = source.artifacts[-1]
            shot_uri = candidate.shot_outputs[0].artifact_uri
            candidate = candidate.model_copy(
                update={
                    "lineage_claims": [
                        (
                            claim.model_copy(
                                update={
                                    "input_artifact_uris": [
                                        *claim.input_artifact_uris,
                                        upstream.uri,
                                    ]
                                }
                            )
                            if claim.artifact_uri == shot_uri
                            else claim
                        )
                        for claim in candidate.lineage_claims
                    ]
                }
            )
            shot_id = delivery.shot_outputs[0].artifact_id
            delivery = delivery.model_copy(
                update={
                    "artifact_graph": delivery.artifact_graph.model_copy(
                        update={
                            "lineages": [
                                (
                                    lineage.model_copy(
                                        update={
                                            "input_artifact_ids": [
                                                *lineage.input_artifact_ids,
                                                upstream.artifact_id,
                                            ]
                                        }
                                    )
                                    if lineage.artifact_id == shot_id
                                    else lineage
                                )
                                for lineage in (
                                    delivery.artifact_graph.lineages
                                )
                            ]
                        }
                    )
                }
            )
            without_context = {
                issue.code
                for issue in evaluate_segment_delivery(
                    package,
                    candidate,
                    delivery,
                )
            }
            self.assertIn(
                "unknown_lineage_artifact_uri",
                without_context,
            )
            verified = load_verified_segment_dependencies(package)
            self.assertEqual(
                evaluate_segment_delivery(
                    package,
                    candidate,
                    delivery,
                    dependency_artifacts=tuple(
                        artifact
                        for dependency in verified
                        for artifact in dependency.delivery.artifacts
                    ),
                ),
                [],
            )

    def test_tool_argument_is_normalized_to_dependency_lineage(self) -> None:
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            package, candidate, _delivery, source, _report_path = (
                self._records(root)
            )
            verified = load_verified_segment_dependencies(package)
            upstream = source.artifacts[-1]
            shot_uri = candidate.shot_outputs[0].artifact_uri
            narration_uri = candidate.narration_artifact_uri
            subtitle_uri = candidate.subtitle_artifact_uri
            primary_uri = candidate.primary_artifact_uri
            records = [
                ToolCallRecord(
                    tool_call_id="make-shot",
                    tool_name="make_video",
                    arguments={
                        "source": package.artifacts[0].uri,
                        "anchor": package.shared_visual_anchor.uri,
                        "previous_segment": upstream.uri,
                    },
                    success=True,
                    output_refs=[shot_uri],
                ),
                ToolCallRecord(
                    tool_call_id="make-narration",
                    tool_name="make_audio",
                    arguments={"source": package.artifacts[0].uri},
                    success=True,
                    output_refs=[narration_uri],
                ),
                ToolCallRecord(
                    tool_call_id="make-subtitle",
                    tool_name="make_subtitle",
                    arguments={"audio": narration_uri},
                    success=True,
                    output_refs=[subtitle_uri],
                ),
                ToolCallRecord(
                    tool_call_id="make-final",
                    tool_name="make_final",
                    arguments={
                        "video": shot_uri,
                        "audio": narration_uri,
                        "subtitle": subtitle_uri,
                    },
                    success=True,
                    output_refs=[primary_uri],
                ),
            ]
            normalized = normalize_candidate_lineage(
                candidate,
                package=package,
                records=records,
                dependency_artifacts=tuple(
                    artifact
                    for dependency in verified
                    for artifact in dependency.delivery.artifacts
                ),
            )
            shot_claim = next(
                item
                for item in normalized.lineage_claims
                if item.artifact_uri == shot_uri
            )
            self.assertIn(upstream.uri, shot_claim.input_artifact_uris)


if __name__ == "__main__":
    unittest.main()