from __future__ import annotations

import hashlib
import json
import os
import shutil
from pathlib import Path
from typing import Any, Callable

from minio import Minio

try:
    import onnxruntime as ort
except ImportError:  # pragma: no cover
    ort = None


class ArtifactInstallError(RuntimeError):
    pass


class ArtifactInstaller:
    def __init__(
        self,
        model_root: Path | None = None,
        registry_path: Path | None = None,
        bucket: str | None = None,
        minio_client: Minio | None = None,
    ) -> None:
        self.model_root = (model_root or Path(os.getenv("RAIL_MODEL_DIR", "/models"))).resolve()
        self.registry_path = registry_path or Path(
            os.getenv("RAIL_MODEL_REGISTRY", "/app/config/model-registry.json")
        )
        self.bucket = bucket or os.getenv("MINIO_ARTIFACT_BUCKET", "rail-model-artifacts")
        self._minio = minio_client

    def status(self) -> dict[str, Any]:
        models: list[dict[str, Any]] = []
        for spec in self._registry().get("models", []):
            expected = self.expected_components(str(spec["model_group"]), str(spec.get("model_version", "")))
            models.append(
                {
                    "model_group": spec["model_group"],
                    "model_version": spec.get("model_version"),
                    "builtin": not bool(spec.get("artifact")),
                    "components": {
                        component: {
                            "path": str(self._target(relative)),
                            "installed": self._target(relative).is_file(),
                        }
                        for component, relative in expected.items()
                    },
                }
            )
        return {
            "status": "ok",
            "profile": self._registry().get("profile"),
            "model_root": str(self.model_root),
            "models": models,
        }

    def expected_components(self, model_group: str, model_version: str) -> dict[str, str]:
        spec = self._model_spec(model_group, model_version)
        if not spec.get("artifact"):
            return {}
        expected = {"primary": str(spec["artifact"])}
        if spec.get("secondary_artifact"):
            expected["secondary"] = str(spec["secondary_artifact"])
        if spec.get("labels"):
            expected["labels"] = str(spec["labels"])
        return expected

    def install(
        self,
        job_id: str,
        model_group: str,
        model_version: str,
        components: list[dict[str, Any]],
        progress: Callable[[str, int, str], None] | None = None,
    ) -> dict[str, Any]:
        notify = progress or (lambda _stage, _progress, _message: None)
        expected = self.expected_components(model_group, model_version)
        if not expected:
            raise ArtifactInstallError("内置运行组件不接受模型制品安装")
        supplied = {str(item.get("component_type")): item for item in components}
        missing = [name for name in expected if name not in supplied]
        if missing:
            raise ArtifactInstallError("缺少必需制品组件: " + "、".join(missing))

        staging = self._target(f".staging/{self._safe_segment(job_id)}")
        backup = staging / "backup"
        shutil.rmtree(staging, ignore_errors=True)
        staging.mkdir(parents=True, exist_ok=True)
        installed: list[dict[str, Any]] = []
        replaced: list[tuple[Path, Path | None]] = []
        try:
            notify("download", 18, "正在从制品库下载模型组件")
            for component_type, relative in expected.items():
                source = supplied[component_type]
                staged_file = staging / "files" / component_type / Path(relative).name
                staged_file.parent.mkdir(parents=True, exist_ok=True)
                self._download(str(source["object_key"]), staged_file)
                actual_sha = self._sha256(staged_file)
                expected_sha = str(source.get("sha256") or "").lower()
                if expected_sha and actual_sha != expected_sha:
                    raise ArtifactInstallError(f"{component_type} 文件 SHA-256 校验失败")
                if component_type in {"primary", "secondary"}:
                    notify("validate", 38, f"正在校验 {component_type} ONNX 制品")
                    model_report = self._validate_onnx(staged_file)
                else:
                    model_report = self._validate_labels(staged_file)
                installed.append(
                    {
                        "component_type": component_type,
                        "relative_path": relative,
                        "sha256": actual_sha,
                        "size_bytes": staged_file.stat().st_size,
                        "validation": model_report,
                    }
                )

            notify("publish", 68, "正在原子发布模型文件")
            for item in installed:
                component_type = str(item["component_type"])
                target = self._target(str(item["relative_path"]))
                staged_file = staging / "files" / component_type / target.name
                target.parent.mkdir(parents=True, exist_ok=True)
                old_backup: Path | None = None
                if target.exists():
                    old_backup = backup / component_type / target.name
                    old_backup.parent.mkdir(parents=True, exist_ok=True)
                    os.replace(target, old_backup)
                os.replace(staged_file, target)
                replaced.append((target, old_backup))

            notify("verify", 86, "正在核对运行目录制品摘要")
            for item in installed:
                target = self._target(str(item["relative_path"]))
                if self._sha256(target) != item["sha256"]:
                    raise ArtifactInstallError(f"发布后摘要不一致: {item['component_type']}")
            notify("completed", 100, "模型制品安装完成")
            return {
                "job_id": job_id,
                "model_group": model_group,
                "model_version": model_version,
                "install_root": str(self.model_root),
                "files": installed,
            }
        except Exception:
            for target, old_backup in reversed(replaced):
                target.unlink(missing_ok=True)
                if old_backup and old_backup.exists():
                    target.parent.mkdir(parents=True, exist_ok=True)
                    os.replace(old_backup, target)
            raise
        finally:
            shutil.rmtree(staging, ignore_errors=True)

    def uninstall(self, model_group: str, model_version: str) -> dict[str, Any]:
        expected = self.expected_components(model_group, model_version)
        if not expected:
            raise ArtifactInstallError("内置运行组件不能卸载")
        removed: list[str] = []
        for relative in expected.values():
            target = self._target(relative)
            if target.is_file():
                target.unlink()
                removed.append(str(target))
            self._remove_empty_parents(target.parent)
        return {"model_group": model_group, "model_version": model_version, "removed": removed}

    def _download(self, object_key: str, target: Path) -> None:
        response = self._client().get_object(self.bucket, self._safe_object_key(object_key))
        try:
            with target.open("wb") as output:
                shutil.copyfileobj(response, output, length=1024 * 1024)
        finally:
            response.close()
            response.release_conn()

    def _client(self) -> Minio:
        if self._minio is None:
            endpoint = os.getenv("MINIO_ENDPOINT", "minio:9000")
            secure = endpoint.startswith("https://")
            endpoint = endpoint.removeprefix("http://").removeprefix("https://")
            self._minio = Minio(
                endpoint,
                access_key=os.getenv("MINIO_ACCESS_KEY", "minioadmin"),
                secret_key=os.getenv("MINIO_SECRET_KEY", "minioadmin"),
                secure=secure,
            )
        return self._minio

    def _registry(self) -> dict[str, Any]:
        if not self.registry_path.is_file():
            raise ArtifactInstallError(f"模型注册表不存在: {self.registry_path}")
        return json.loads(self.registry_path.read_text(encoding="utf-8"))

    def _model_spec(self, model_group: str, model_version: str) -> dict[str, Any]:
        group_fallback: dict[str, Any] | None = None
        for spec in self._registry().get("models", []):
            if str(spec.get("model_group")) != model_group:
                continue
            group_fallback = spec
            if str(spec.get("model_version")) == model_version:
                return spec
        if group_fallback and not model_version:
            return group_fallback
        raise ArtifactInstallError(f"运行注册表未配置模型 {model_group}:{model_version}")

    def _target(self, relative: str) -> Path:
        target = (self.model_root / relative).resolve()
        if target != self.model_root and self.model_root not in target.parents:
            raise ArtifactInstallError("模型安装路径越界")
        return target

    def _validate_onnx(self, path: Path) -> dict[str, Any]:
        if path.suffix.lower() != ".onnx":
            raise ArtifactInstallError("模型组件必须是 .onnx 文件")
        if ort is None:
            raise ArtifactInstallError("安装代理未安装 ONNX Runtime")
        try:
            session = ort.InferenceSession(str(path), providers=["CPUExecutionProvider"])
            return {
                "inputs": [self._node_meta(item) for item in session.get_inputs()],
                "outputs": [self._node_meta(item) for item in session.get_outputs()],
                "providers": session.get_providers(),
            }
        except Exception as exc:
            raise ArtifactInstallError(f"ONNX 制品无法创建推理会话: {type(exc).__name__}: {exc}") from exc

    def _validate_labels(self, path: Path) -> dict[str, Any]:
        try:
            labels = [line.strip() for line in path.read_text(encoding="utf-8-sig").splitlines() if line.strip()]
        except Exception as exc:
            raise ArtifactInstallError(f"标签文件不是有效 UTF-8 文本: {exc}") from exc
        if not labels:
            raise ArtifactInstallError("标签文件不能为空")
        if len(labels) != len(set(labels)):
            raise ArtifactInstallError("标签文件包含重复类别")
        return {"label_count": len(labels)}

    @staticmethod
    def _node_meta(node: Any) -> dict[str, Any]:
        return {"name": node.name, "type": node.type, "shape": [str(value) for value in node.shape]}

    @staticmethod
    def _sha256(path: Path) -> str:
        digest = hashlib.sha256()
        with path.open("rb") as stream:
            for block in iter(lambda: stream.read(1024 * 1024), b""):
                digest.update(block)
        return digest.hexdigest()

    @staticmethod
    def _safe_segment(value: str) -> str:
        if not value or not all(char.isalnum() or char in "-_" for char in value):
            raise ArtifactInstallError("非法安装任务编号")
        return value

    @staticmethod
    def _safe_object_key(value: str) -> str:
        normalized = value.replace("\\", "/").lstrip("/")
        if not normalized or "../" in normalized or normalized == "..":
            raise ArtifactInstallError("非法对象存储路径")
        return normalized

    def _remove_empty_parents(self, path: Path) -> None:
        current = path
        while current != self.model_root and self.model_root in current.parents:
            try:
                current.rmdir()
            except OSError:
                break
            current = current.parent
