from __future__ import annotations

import hashlib
import json
import zipfile
from pathlib import Path
from typing import Any

import pytest
from PIL import Image, ImageDraw
from test_perception_qualification import _evaluation_fixture

from k1link.compute import (
    AnnotationWorkspaceError,
    prepare_annotation_workspace,
    prepare_recorded_evaluation_pack,
    validate_annotation_workspace,
)
from k1link.device_plugins.xgrids_k1.analyze.valid_fov import validate_k1_valid_fov_mask


def _canonical_json(value: object) -> bytes:
    return json.dumps(
        value,
        ensure_ascii=False,
        sort_keys=True,
        separators=(",", ":"),
        allow_nan=False,
    ).encode("utf-8")


def _sha256(path: Path) -> str:
    return hashlib.sha256(path.read_bytes()).hexdigest()


def _write_json(path: Path, payload: object) -> None:
    path.write_text(json.dumps(payload, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")


def _artifact(path: Path, root: Path) -> dict[str, Any]:
    return {
        "path": path.relative_to(root).as_posix(),
        "byte_length": path.stat().st_size,
        "sha256": _sha256(path),
    }


def _prelabel_fixture(pack_root: Path, valid_fov_root: Path, output_root: Path) -> Path:
    pack = json.loads((pack_root / "manifest.json").read_text(encoding="utf-8"))
    valid_fov = validate_k1_valid_fov_mask(valid_fov_root)
    staging = output_root / "staging"
    (staging / "instance-prelabels").mkdir(parents=True)
    (staging / "semantic-prelabels").mkdir()
    with Image.open(valid_fov.mask_path) as opened:
        valid_mask = opened.copy()
    rows: list[dict[str, Any]] = []
    for order, source in enumerate(pack["identity"]["frames"], start=1):
        instance = Image.new("I;16", (valid_fov.width, valid_fov.height), 0)
        ImageDraw.Draw(instance).rectangle((390, 290, 409, 309), fill=1)
        instance.save(staging / "instance-prelabels" / f"image-{order:03d}.png")
        semantic = Image.composite(
            Image.new("L", (valid_fov.width, valid_fov.height), 7),
            Image.new("L", (valid_fov.width, valid_fov.height), 0),
            valid_mask,
        )
        ImageDraw.Draw(semantic).rectangle((390, 290, 409, 309), fill=4)
        semantic.save(staging / "semantic-prelabels" / f"image-{order:03d}.png")
        rows.append(
            {
                "schema_version": "missioncore.perception-evaluation-prelabel-frame/v1",
                "image_id": source["image_id"],
                "frame_index": source["frame_index"],
                "session_seconds": source["session_seconds"],
                "role": source["role"],
                "group_id": source["group_id"],
                "instances": [
                    {
                        "instance_id": 1,
                        "draft_category_id": 4,
                        "draft_category": "car",
                        "source_model_category_id": 3,
                        "source_model_category": "car",
                        "score": 0.9,
                        "box_xyxy": [390.0, 290.0, 410.0, 310.0],
                        "mask_pixels": 400,
                        "review_state": "unreviewed-model-draft",
                    }
                ],
                "review_state": "unreviewed-model-draft",
            }
        )
    (staging / "frames.jsonl").write_text(
        "".join(json.dumps(row, separators=(",", ":")) + "\n" for row in rows),
        encoding="utf-8",
    )
    identity = {
        "schema_version": "missioncore.perception-evaluation-prelabels-identity/v1",
        "evaluation_pack_id": pack["generation_id"],
        "evaluation_identity_sha256": pack["identity_sha256"],
        "valid_fov_generation_id": valid_fov.generation_id,
        "pipeline": "fixture/v1",
    }
    identity_sha256 = hashlib.sha256(_canonical_json(identity)).hexdigest()
    result_id = f"evaluation-prelabels-{identity_sha256}"
    artifacts = [
        _artifact(path, staging) for path in sorted(staging.rglob("*")) if path.is_file()
    ]
    _write_json(
        staging / "result.json",
        {
            "schema_version": "missioncore.perception-evaluation-prelabels/v1",
            "result_id": result_id,
            "identity_sha256": identity_sha256,
            "identity": identity,
            "review_state": "unreviewed-model-draft",
            "artifacts": artifacts,
        },
    )
    final = output_root / result_id
    staging.rename(final)
    return final


def _workspace_fixture(tmp_path: Path) -> tuple[Path, Path, Path, Path]:
    job_root, qualification_root, valid_fov_root, frames_root, timeline_path, requests = (
        _evaluation_fixture(tmp_path)
    )
    pack = prepare_recorded_evaluation_pack(
        job_root=job_root,
        qualification_root=qualification_root,
        valid_fov_root=valid_fov_root,
        decoded_frames_root=frames_root,
        timeline_path=timeline_path,
        output_root=tmp_path / "evaluation-packs",
        selection=requests,
        decoder_version="fixture-decoder/v1",
        selection_document_sha256="1" * 64,
        producer_files=(("fixture.py", "2" * 64),),
    )
    prelabels = _prelabel_fixture(pack.root, valid_fov_root, tmp_path / "prelabels")
    return pack.root, prelabels, valid_fov_root, tmp_path / "annotation-workspaces"


def test_annotation_workspace_is_reproducible_and_remains_unreviewed(tmp_path: Path) -> None:
    pack_root, prelabels_root, valid_fov_root, output_root = _workspace_fixture(tmp_path)
    first = prepare_annotation_workspace(
        evaluation_pack_root=pack_root,
        prelabels_root=prelabels_root,
        valid_fov_root=valid_fov_root,
        output_root=output_root,
        producer_files=(("fixture.py", "3" * 64),),
    )
    repeated = prepare_annotation_workspace(
        evaluation_pack_root=pack_root,
        prelabels_root=prelabels_root,
        valid_fov_root=valid_fov_root,
        output_root=output_root,
        producer_files=(("fixture.py", "3" * 64),),
    )

    assert repeated == first
    assert first.frame_count == 22
    assert first.draft_instance_count == 22
    manifest = json.loads(first.manifest_path.read_text(encoding="utf-8"))
    assert manifest["ground_truth"] is False
    assert manifest["state"] == "prepared-unreviewed-model-draft"
    review = json.loads(first.review_template_path.read_text(encoding="utf-8"))
    assert {row["review_status"] for row in review["frames"]} == {"unreviewed"}

    with zipfile.ZipFile(first.instance_archive_path) as archive:
        coco = json.loads(archive.read("annotations/instances_default.json"))
    assert len(coco["images"]) == 22
    assert len(coco["annotations"]) == 22
    assert all(sum(row["segmentation"]["counts"]) == 800 * 600 for row in coco["annotations"])
    assert all(row["area"] == 400 for row in coco["annotations"])
    assert (
        validate_annotation_workspace(
            first.root,
            evaluation_pack_root=pack_root,
            prelabels_root=prelabels_root,
            valid_fov_root=valid_fov_root,
        )
        == first
    )


def test_annotation_workspace_rejects_changed_archive(tmp_path: Path) -> None:
    pack_root, prelabels_root, valid_fov_root, output_root = _workspace_fixture(tmp_path)
    result = prepare_annotation_workspace(
        evaluation_pack_root=pack_root,
        prelabels_root=prelabels_root,
        valid_fov_root=valid_fov_root,
        output_root=output_root,
        producer_files=(("fixture.py", "3" * 64),),
    )
    result.semantic_archive_path.write_bytes(b"changed")

    with pytest.raises(AnnotationWorkspaceError, match="artifact changed"):
        validate_annotation_workspace(result.root)
