import json
import hashlib
import io
from pathlib import Path

import onnx
import pytest
from onnx import TensorProto, helper

from app.installer import ArtifactInstallError, ArtifactInstaller


class ObjectResponse(io.BytesIO):
    def release_conn(self) -> None:
        return None


class ObjectStore:
    def __init__(self, objects: dict[str, bytes]) -> None:
        self.objects = objects

    def get_object(self, _bucket: str, object_key: str) -> ObjectResponse:
        return ObjectResponse(self.objects[object_key])


def identity_onnx() -> bytes:
    input_info = helper.make_tensor_value_info("input", TensorProto.FLOAT, [1, 3])
    output_info = helper.make_tensor_value_info("output", TensorProto.FLOAT, [1, 3])
    graph = helper.make_graph([helper.make_node("Identity", ["input"], ["output"])], "identity", [input_info], [output_info])
    model = helper.make_model(graph, producer_name="artifact-installer-test", opset_imports=[helper.make_opsetid("", 13)])
    model.ir_version = 10
    return model.SerializeToString()


def registry(path: Path) -> Path:
    payload = {
        "profile": "cpu-local",
        "models": [
            {
                "model_group": "thermal-analyzer",
                "model_version": "cpu-v1.0.0",
                "artifact": "thermal/main.onnx",
                "secondary_artifact": "thermal/patchcore.onnx",
                "labels": "thermal/labels.txt",
            },
            {
                "model_group": "pointcloud-analyzer",
                "model_version": "cpu-v1.0.0",
                "artifact": None,
            },
        ],
    }
    path.write_text(json.dumps(payload), encoding="utf-8")
    return path


def test_expected_components_and_builtin(tmp_path: Path):
    service = ArtifactInstaller(tmp_path / "models", registry(tmp_path / "registry.json"))
    assert set(service.expected_components("thermal-analyzer", "cpu-v1.0.0")) == {
        "primary",
        "secondary",
        "labels",
    }
    assert service.expected_components("pointcloud-analyzer", "cpu-v1.0.0") == {}


def test_install_rejects_incomplete_composite_package(tmp_path: Path):
    service = ArtifactInstaller(tmp_path / "models", registry(tmp_path / "registry.json"))
    with pytest.raises(ArtifactInstallError, match="secondary"):
        service.install(
            "job-1",
            "thermal-analyzer",
            "cpu-v1.0.0",
            [{"component_type": "primary", "object_key": "a", "sha256": "", "size_bytes": 1}],
        )


def test_uninstall_removes_only_registered_paths(tmp_path: Path):
    model_root = tmp_path / "models"
    service = ArtifactInstaller(model_root, registry(tmp_path / "registry.json"))
    for relative in service.expected_components("thermal-analyzer", "cpu-v1.0.0").values():
        target = model_root / relative
        target.parent.mkdir(parents=True, exist_ok=True)
        target.write_bytes(b"content")
    keep = model_root / "other" / "keep.onnx"
    keep.parent.mkdir(parents=True)
    keep.write_bytes(b"keep")

    result = service.uninstall("thermal-analyzer", "cpu-v1.0.0")

    assert len(result["removed"]) == 3
    assert keep.is_file()


def test_install_validates_and_atomically_publishes_package(tmp_path: Path):
    model = identity_onnx()
    labels = "normal\ndefect\n".encode()
    objects = {"main": model, "secondary": model, "labels": labels}
    service = ArtifactInstaller(
        tmp_path / "models",
        registry(tmp_path / "registry.json"),
        minio_client=ObjectStore(objects),
    )
    components = [
        {
            "component_type": component,
            "object_key": object_key,
            "sha256": hashlib.sha256(objects[object_key]).hexdigest(),
            "size_bytes": len(objects[object_key]),
        }
        for component, object_key in (("primary", "main"), ("secondary", "secondary"), ("labels", "labels"))
    ]

    result = service.install("job-valid", "thermal-analyzer", "cpu-v1.0.0", components)

    assert len(result["files"]) == 3
    for relative in service.expected_components("thermal-analyzer", "cpu-v1.0.0").values():
        assert (tmp_path / "models" / relative).is_file()
