from __future__ import annotations

import json
import tempfile
import unittest
from pathlib import Path

from agents.semantic_evaluator import SemanticEvaluator
from pipelines.evaluate_semantic_predictions import run_semantic_evaluation
from pipelines.prepare_validation_reference import run_validation_reference_build
from pipelines.run_validation_inference import ValidationInferencePipeline


class SemanticEvaluatorTest(unittest.TestCase):
    def test_evaluate_predictions_matches_targets_and_reports_metrics(self) -> None:
        with tempfile.TemporaryDirectory() as temp_dir:
            temp_path = Path(temp_dir)
            reference_path = temp_path / "reference.jsonl"
            prediction_path = temp_path / "prediction.jsonl"

            reference_rows = [
                {
                    "image_path": "/tmp/img_a.jpg",
                    "annotations": [
                        {
                            "bbox_1000": [100, 100, 200, 200],
                            "maturity_level": "未成熟",
                            "maturity_ratio": 0.1,
                            "occlusion_degree": "无",
                            "reasoning": "ref1",
                        },
                        {
                            "bbox_1000": [300, 300, 400, 400],
                            "maturity_level": "完熟",
                            "maturity_ratio": 0.9,
                            "occlusion_degree": "轻度",
                            "reasoning": "ref2",
                        },
                    ],
                }
            ]
            prediction_rows = [
                {
                    "image_path": "/workspace/img_a.jpg",
                    "annotations": [
                        {
                            "bbox_1000": [102, 102, 198, 198],
                            "maturity_level": "未成熟",
                            "maturity_ratio": 0.12,
                            "occlusion_degree": "无",
                            "reasoning": "pred1",
                        },
                        {
                            "bbox_1000": [302, 302, 398, 398],
                            "maturity_level": "半成熟",
                            "maturity_ratio": 0.5,
                            "occlusion_degree": "轻度",
                            "reasoning": "pred2",
                        },
                        {
                            "bbox_1000": [700, 700, 800, 800],
                            "maturity_level": "未成熟",
                            "maturity_ratio": 0.1,
                            "occlusion_degree": "无",
                            "reasoning": "extra",
                        },
                    ],
                }
            ]

            reference_path.write_text("\n".join(json.dumps(row, ensure_ascii=False) for row in reference_rows) + "\n", encoding="utf-8")
            prediction_path.write_text("\n".join(json.dumps(row, ensure_ascii=False) for row in prediction_rows) + "\n", encoding="utf-8")

            evaluator = SemanticEvaluator(config={"match_iou_threshold": 0.5, "risk_top_k": 5, "bbox_coordinate_space": "bbox_1000"})
            result = evaluator.evaluate_predictions(
                reference_dataset_path=str(reference_path),
                prediction_dataset_path=str(prediction_path),
            )

            self.assertEqual(result["summary"]["image_count"], 1)
            self.assertEqual(result["summary"]["reference_annotations_total"], 2)
            self.assertEqual(result["summary"]["prediction_annotations_total"], 3)
            self.assertEqual(result["summary"]["matched_annotations_total"], 2)
            self.assertAlmostEqual(result["summary"]["precision"], 0.6667, places=4)
            self.assertAlmostEqual(result["summary"]["recall"], 1.0, places=4)
            self.assertAlmostEqual(result["summary"]["f1"], 0.8, places=4)
            self.assertAlmostEqual(result["summary"]["maturity_accuracy_on_matched"], 0.5, places=4)

            per_image = result["per_image_results"][0]
            self.assertEqual(per_image["image_key"], "img_a.jpg")
            self.assertEqual(per_image["matched_count"], 2)
            self.assertEqual(per_image["unmatched_prediction_count"], 1)
            self.assertAlmostEqual(per_image["f1"], 0.8, places=4)

    def test_evaluator_can_compare_pixel_predictions_against_bbox_1000_reference(self) -> None:
        with tempfile.TemporaryDirectory() as temp_dir:
            temp_path = Path(temp_dir)
            reference_path = temp_path / "reference.jsonl"
            prediction_path = temp_path / "prediction.jsonl"

            reference_rows = [
                {
                    "image_path": "/tmp/img_c.jpg",
                    "annotations": [
                        {
                            "bbox_1000": [100, 100, 300, 300],
                            "image_width": 1000,
                            "image_height": 500,
                        }
                    ],
                }
            ]
            prediction_rows = [
                {
                    "image_path": "/tmp/img_c.jpg",
                    "annotations": [
                        {
                            "bbox": [100, 50, 300, 150],
                            "image_width": 1000,
                            "image_height": 500,
                        }
                    ],
                }
            ]

            reference_path.write_text(json.dumps(reference_rows[0], ensure_ascii=False) + "\n", encoding="utf-8")
            prediction_path.write_text(json.dumps(prediction_rows[0], ensure_ascii=False) + "\n", encoding="utf-8")

            evaluator = SemanticEvaluator(config={"bbox_coordinate_space": "bbox_1000"})
            result = evaluator.evaluate_predictions(str(reference_path), str(prediction_path))
            self.assertEqual(result["summary"]["matched_annotations_total"], 1)
            self.assertAlmostEqual(result["summary"]["precision"], 1.0, places=4)
            self.assertAlmostEqual(result["summary"]["recall"], 1.0, places=4)

    def test_pipeline_writes_summary_and_jsonl_outputs(self) -> None:
        with tempfile.TemporaryDirectory() as temp_dir:
            temp_path = Path(temp_dir)
            reference_path = temp_path / "reference.jsonl"
            prediction_path = temp_path / "prediction.jsonl"
            output_dir = temp_path / "semantic_eval"

            reference_gold_row = {
                "image_path": "/tmp/img_b.jpg",
                "source_records": [
                    {
                        "bbox": [40, 50, 80, 100],
                        "bbox_1000": [100, 100, 200, 200],
                        "image_width": 400,
                        "image_height": 500,
                        "maturity_level": "未成熟",
                        "maturity_ratio": 0.1,
                        "occlusion_degree": "无",
                        "reasoning": "gold",
                    }
                ],
                "messages": [
                    {"role": "assistant", "content": json.dumps({"annotations": [{"bbox": [100, 100, 200, 200]}]}, ensure_ascii=False)}
                ],
            }
            prediction_row = {
                "image_path": "/tmp/img_b.jpg",
                "annotations": [
                    {
                        "bbox_1000": [100, 100, 200, 200],
                        "maturity_level": "未成熟",
                        "maturity_ratio": 0.1,
                        "occlusion_degree": "无",
                        "reasoning": "pred",
                    }
                ],
            }

            reference_path.write_text(json.dumps(reference_gold_row, ensure_ascii=False) + "\n", encoding="utf-8")
            prediction_path.write_text(json.dumps(prediction_row, ensure_ascii=False) + "\n", encoding="utf-8")

            summary = run_semantic_evaluation(
                reference_dataset_path=reference_path,
                prediction_dataset_path=prediction_path,
                output_dir=output_dir,
                bbox_coordinate_space="bbox_1000",
            )

            self.assertTrue((output_dir / "summary.json").exists())
            self.assertTrue((output_dir / "report.md").exists())
            self.assertTrue((output_dir / "per_image_results.jsonl").exists())
            self.assertTrue((output_dir / "risk_cases.jsonl").exists())
            self.assertEqual(summary["matched_annotations_total"], 1)
            self.assertEqual(summary["bbox_coordinate_space"], "bbox_1000")

    def test_validation_reference_builder_keeps_only_val_images(self) -> None:
        with tempfile.TemporaryDirectory() as temp_dir:
            temp_path = Path(temp_dir)
            gold_path = temp_path / "gold.jsonl"
            val_path = temp_path / "val.jsonl"
            output_dir = temp_path / "reference"

            gold_rows = [
                {"image_path": "/tmp/a.jpg", "source_records": []},
                {"image_path": "/tmp/b.jpg", "source_records": []},
            ]
            val_rows = [
                {"images": ["/tmp/b.jpg"], "metadata": {"image_path": "/tmp/b.jpg"}},
            ]
            gold_path.write_text("\n".join(json.dumps(row, ensure_ascii=False) for row in gold_rows) + "\n", encoding="utf-8")
            val_path.write_text(json.dumps(val_rows[0], ensure_ascii=False) + "\n", encoding="utf-8")

            summary = run_validation_reference_build(gold_path, val_path, output_dir)
            subset_lines = (output_dir / "reference_subset.jsonl").read_text(encoding="utf-8").splitlines()
            self.assertEqual(summary["matched_reference_rows"], 1)
            self.assertEqual(len(subset_lines), 1)
            self.assertIn("/tmp/b.jpg", subset_lines[0])

    def test_validation_inference_parser_converts_bbox_1000_to_pixel(self) -> None:
        pipeline = ValidationInferencePipeline()
        annotations, error = pipeline._parse_prediction_annotations(
            raw_text='[{"bbox_2d": [100, 200, 300, 400], "label": "未成熟番茄"}]',
            image_width=1000,
            image_height=500,
        )
        self.assertIsNone(error)
        self.assertEqual(len(annotations), 1)
        self.assertEqual(annotations[0]["bbox_1000"], [100.0, 200.0, 300.0, 400.0])
        self.assertEqual(annotations[0]["bbox"], [100.0, 100.0, 300.0, 200.0])
        self.assertEqual(annotations[0]["maturity_level"], "未成熟")

    def test_validation_inference_parser_accepts_string_bbox_inside_think(self) -> None:
        pipeline = ValidationInferencePipeline()
        raw_text = (
            '<think>\n'
            '[{"bbox_2d": "[290, 526, 380, 754]", "label": "未成熟番茄"}]\n'
            '</think>\n\n'
            '[{"bbox_2d": "[290, 526, 380, 754]", "label": "未成熟番茄"}]'
        )
        annotations, error = pipeline._parse_prediction_annotations(
            raw_text=raw_text,
            image_width=4480,
            image_height=2016,
        )
        self.assertIsNone(error)
        self.assertEqual(len(annotations), 1)
        self.assertEqual(annotations[0]["bbox_1000"], [290.0, 526.0, 380.0, 754.0])
        self.assertEqual(annotations[0]["maturity_level"], "未成熟")

    def test_validation_inference_parser_ignores_bbox_lists_in_reasoning(self) -> None:
        pipeline = ValidationInferencePipeline()
        raw_text = (
            '<think>\n'
            '候选目标位于 [290, 526, 380, 754]，另一个位于 [452, 703, 544, 926]。\n'
            '</think>\n\n'
            '[{"bbox_2d": "[452, 703, 544, 926]", "label": "未成熟番茄", '
            '"attribute_text": "无遮挡,局部可见绿色"}]'
        )
        annotations, error = pipeline._parse_prediction_annotations(
            raw_text=raw_text,
            image_width=4480,
            image_height=2016,
        )
        self.assertIsNone(error)
        self.assertEqual(len(annotations), 1)
        self.assertEqual(annotations[0]["bbox_1000"], [452.0, 703.0, 544.0, 926.0])
        self.assertEqual(annotations[0]["occlusion_degree"], "无")

    def test_validation_inference_request_keeps_output_schema_as_user_constraint(self) -> None:
        class DummyInferRequest:
            def __init__(self, messages: list[dict[str, str]], images: list[str]) -> None:
                self.messages = messages
                self.images = images

        pipeline = ValidationInferencePipeline()
        row = {
            "messages": [
                {"role": "user", "content": "<image>找到图像中的番茄。"},
                {
                    "role": "assistant",
                    "content": '[{"bbox_2d": "<bbox>", "label": "<ref-object>", "attribute_text": "无遮挡,局部可见绿色"}]',
                },
            ],
            "images": ["/tmp/demo.jpg"],
        }

        request = pipeline._build_infer_request(row, DummyInferRequest)

        self.assertEqual(len(request.messages), 2)
        self.assertEqual(request.messages[0]["role"], "user")
        self.assertEqual(request.messages[0]["content"], "<image>找到图像中的番茄。")
        self.assertEqual(request.messages[1]["role"], "user")
        self.assertIn("严格按照下面的 JSON 模板输出最终答案", request.messages[1]["content"])
        self.assertIn('"attribute_text"', request.messages[1]["content"])
        self.assertEqual(request.images, ["/tmp/demo.jpg"])

    def test_validation_inference_can_override_output_template_for_ablation(self) -> None:
        pipeline = ValidationInferencePipeline(output_template_mode="official")
        prompt = pipeline._build_output_format_prompt(
            '[{"bbox_2d": "<bbox>", "label": "<ref-object>", "attribute_text": "无遮挡,局部可见绿色"}]'
        )

        self.assertIn('"bbox_2d"', prompt)
        self.assertIn('"label"', prompt)
        self.assertNotIn("attribute_text", prompt)

    def test_validation_inference_parser_accepts_structured_occlusion_field(self) -> None:
        pipeline = ValidationInferencePipeline()
        annotations, parse_error = pipeline._parse_prediction_annotations(
            raw_text='[{"bbox_2d": "[100, 200, 300, 400]", "label": "未成熟番茄", '
            '"occlusion_degree": "轻度", "visible_color": "绿色"}]',
            image_width=1000,
            image_height=1000,
        )

        self.assertIsNone(parse_error)
        self.assertEqual(annotations[0]["occlusion_degree"], "轻度")
        self.assertEqual(annotations[0]["visible_color"], "绿色")


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