from __future__ import annotations

import json
import tempfile
import unittest
from pathlib import Path

from cadfs_to_cdsl.describe import describe_samples, select_description_samples


SAMPLE_FS = r'''
FeatureScript 1511;
import(path : "onshape/std/geometry.fs", version : "1511.0");
const mm = millimeter;
const FACE = EntityType.FACE;
function v(x, y){return vector(x, y);}
function sQuery(a, b, c) {return sketchEntityQuery(a, b, c);}
export const myFeature = defineFeature(function(context is Context, id is Id, definition is map)
    precondition{}
    {
        {
            var Q0;
            Q0=qCreatedBy(makeId("Top.planeOp"),FACE);
            var sketch = newSketch(context, id + "F0", { "sketchPlane" : qUnion([Q0])});
            skLineSegment(sketch, "E0", {"start": v(-269.41, -156.6) * mm, "end": v(-0.92, 311.62) * mm});
            skLineSegment(sketch, "E1", {"start": v(-0.92, 311.62) * mm, "end": v(270.33, -155.02) * mm});
            skLineSegment(sketch, "E2", {"start": v(270.33, -155.02) * mm, "end": v(-269.41, -156.6) * mm});
            skSolve(sketch);
        }
        {
            var Q0;
            Q0 = qSketchRegion(id + "F0", true);
            extrude(context, id + "F1", {"entities" : qUnion([Q0]), "depth" : 1828.8 * mm});
        }
    });
'''


SAMPLE_ANNOTATION = """Step 1 - Sketch
Draw a closed triangle.

Step 2 - Extrude NEW
Extrude the triangular area upward a distance of 1828.8 mm.
"""


class FailingVision:
    def describe(self, *, sample_id: str, image_path: Path, local_facts: dict[str, object]) -> dict[str, object]:
        raise RuntimeError("vision offline")


class StaticVision:
    def describe(self, *, sample_id: str, image_path: Path, local_facts: dict[str, object]) -> dict[str, object]:
        return {
            "category": "视觉增强三角楔块（候选）",
            "category_confidence": 0.7,
            "candidate_names": ["视觉楔块"],
            "summary_zh": "视觉模型认为该模型是三角楔块候选件。",
            "possible_functions": ["可能用于定位。"],
            "applications": ["夹具。"],
            "structural_features": ["斜面"],
            "geometric_features": ["三角截面"],
            "keywords_zh": ["视觉识别"],
            "keywords_en": ["vision wedge"],
            "uncertainties": ["视觉语义仍需人工确认。"],
        }


class DescribeTests(unittest.TestCase):
    def test_selects_only_requested_shard(self):
        with tempfile.TemporaryDirectory() as tmp:
            root = Path(tmp)
            self._write_sample(root, "0005", "00050089")
            self._write_sample(root, "0006", "00060001")

            selected = select_description_samples(root, shard="0005")

            self.assertEqual([sample.sample_id for sample in selected], ["00050089"])

    def test_hybrid_falls_back_and_writes_txt_manifest(self):
        with tempfile.TemporaryDirectory() as tmp:
            root = Path(tmp)
            self._write_sample(root, "0005", "00050089")

            records, summary = describe_samples(root, shard="0005", mode="hybrid", vision_client=FailingVision())

            self.assertEqual(summary["sample_count"], 1)
            self.assertEqual(summary["statuses"], {"local_fallback": 1})
            txt_path = root / "description_txt/0005/00050089.txt"
            self.assertTrue(txt_path.is_file())
            text = txt_path.read_text(encoding="utf-8")
            self.assertIn("样本ID：00050089", text)
            self.assertIn("三角棱柱", text)
            self.assertIn("证据与不确定性", text)
            manifest = root / "description_txt/description_manifest.jsonl"
            rows = [json.loads(line) for line in manifest.read_text(encoding="utf-8").splitlines()]
            self.assertEqual(len(rows), 1)
            self.assertEqual(rows[0]["sample_id"], "00050089")
            self.assertEqual(rows[0]["status"], "local_fallback")
            self.assertIn("vision_failed", [item["code"] for item in records[0]["diagnostics"]])

    def test_hybrid_uses_vision_when_available(self):
        with tempfile.TemporaryDirectory() as tmp:
            root = Path(tmp)
            self._write_sample(root, "0005", "00050089")

            records, _summary = describe_samples(root, shard="0005", mode="hybrid", vision_client=StaticVision())

            self.assertEqual(records[0]["status"], "described_hybrid")
            self.assertEqual(records[0]["category"], "视觉增强三角楔块（候选）")
            self.assertIn("视觉楔块", records[0]["candidate_names"])

    def test_existing_txt_is_skipped_without_force(self):
        with tempfile.TemporaryDirectory() as tmp:
            root = Path(tmp)
            self._write_sample(root, "0005", "00050089")
            first, _ = describe_samples(root, shard="0005", mode="local")
            txt_path = root / "description_txt/0005/00050089.txt"
            before = txt_path.read_text(encoding="utf-8")

            second, _ = describe_samples(root, shard="0005", mode="local")

            self.assertEqual(second[0]["status"], first[0]["status"])
            self.assertEqual(txt_path.read_text(encoding="utf-8"), before)

    @staticmethod
    def _write_sample(root: Path, shard: str, sample_id: str) -> None:
        for directory, suffix, content in (
            ("featurescript_rp", ".txt", SAMPLE_FS),
            ("text_annotations", ".txt", SAMPLE_ANNOTATION),
            ("step_abc", ".step", "ISO-10303-21;\nEND-ISO-10303-21;\n"),
            ("stl_abc", ".stl", "solid sample\nendsolid sample\n"),
        ):
            path = root / directory / shard / f"{sample_id}{suffix}"
            path.parent.mkdir(parents=True, exist_ok=True)
            path.write_text(content, encoding="utf-8")
        image = root / "multiview_images_abc" / shard / f"{sample_id}.png"
        image.parent.mkdir(parents=True, exist_ok=True)
        image.write_bytes(
            b"\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR"
            b"\x00\x00\x00\x01\x00\x00\x00\x01\x08\x02\x00\x00\x00"
            b"\x90wS\xde\x00\x00\x00\x00IEND\xaeB`\x82"
        )


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