from __future__ import annotations

import hashlib
import json
import os
import shutil
import subprocess
from datetime import datetime, timezone
from pathlib import Path
from typing import Iterable

from ...config import settings

WEIGHT_EXTS = {".pt", ".pth", ".onnx", ".engine"}


def sha256_file(path: Path, chunk_size: int = 1024 * 1024 * 8) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        while True:
            chunk = handle.read(chunk_size)
            if not chunk:
                break
            digest.update(chunk)
    return digest.hexdigest()


def iter_source_weights() -> Iterable[Path]:
    project_root = settings.project_root.resolve()
    for path in settings.source_root.rglob("*"):
        if not path.is_file() or path.suffix.lower() not in WEIGHT_EXTS:
            continue
        try:
            path.resolve().relative_to(project_root)
            continue
        except ValueError:
            yield path


def classify_weight(path: Path) -> dict[str, str]:
    try:
        rel = str(path.resolve().relative_to(settings.source_root))
    except ValueError:
        rel = str(path)
    lower = rel.lower()
    if "yolo" in lower:
        family = "yolo"
    elif "mmseg" in lower or "my_local_model" in lower:
        family = "mmseg"
    elif "segmodel" in lower:
        family = "segmodel"
    else:
        family = "misc"
    if "best.pt" in lower or "best.pth" in lower:
        role = "trained_best"
    elif "last.pt" in lower or "last.pth" in lower:
        role = "trained_last"
    elif "pretrain" in lower or "my_local_model" in lower:
        role = "pretrained"
    else:
        role = "weight"
    return {"family": family, "role": role}


def copy_weight(src: Path, dst: Path, mode: str) -> None:
    dst.parent.mkdir(parents=True, exist_ok=True)
    if mode == "hardlink":
        if dst.exists():
            dst.unlink()
        os.link(src, dst)
    elif mode == "reflink":
        subprocess.run(["cp", "--reflink=auto", "--preserve=timestamps", str(src), str(dst)], check=True)
    else:
        shutil.copy2(src, dst)


def sync_weights(mode: str = "copy", hash_files: bool = True, skip_existing: bool = True) -> dict:
    files_dir = settings.weights_root / "files"
    settings.weights_root.mkdir(parents=True, exist_ok=True)
    entries = []
    total_bytes = 0
    for src in sorted(iter_source_weights()):
        rel = src.resolve().relative_to(settings.source_root)
        dst = files_dir / rel
        stat = src.stat()
        total_bytes += stat.st_size
        copied = False
        if not (skip_existing and dst.exists() and dst.stat().st_size == stat.st_size):
            copy_weight(src, dst, mode)
            copied = True
        meta = classify_weight(src)
        entry = {
            "source_path": str(rel),
            "stored_path": str(dst.resolve().relative_to(settings.project_root)),
            "size": stat.st_size,
            "family": meta["family"],
            "role": meta["role"],
            "copied": copied,
        }
        if hash_files:
            entry["sha256"] = sha256_file(dst)
        entries.append(entry)

    manifest = {
        "generated_at": datetime.now(timezone.utc).isoformat(),
        "updated_at": datetime.now(timezone.utc).isoformat(),
        "source_root": str(settings.source_root),
        "mode": mode,
        "count": len(entries),
        "total_bytes": total_bytes,
        "files": entries,
    }
    manifest_path = settings.weights_root / "manifest.json"
    manifest_path.write_text(json.dumps(manifest, ensure_ascii=False, indent=2), encoding="utf-8")
    return manifest


def load_manifest() -> dict:
    manifest_path = settings.weights_root / "manifest.json"
    if not manifest_path.exists():
        return {
            "generated_at": None,
            "source_root": str(settings.source_root),
            "count": 0,
            "total_bytes": 0,
            "files": [],
        }
    return json.loads(manifest_path.read_text(encoding="utf-8"))


def verify_weights() -> dict:
    manifest = load_manifest()
    checked = []
    ok_count = 0
    for entry in manifest.get("files", []):
        path = settings.project_root / entry["stored_path"]
        exists = path.exists()
        size_ok = exists and path.stat().st_size == entry.get("size")
        hash_ok = None
        if exists and "sha256" in entry:
            hash_ok = sha256_file(path) == entry["sha256"]
        ok = bool(exists and size_ok and (hash_ok is not False))
        ok_count += int(ok)
        checked.append(
            {
                "stored_path": entry["stored_path"],
                "exists": exists,
                "size_ok": size_ok,
                "hash_ok": hash_ok,
                "ok": ok,
            }
        )
    return {"count": len(checked), "ok_count": ok_count, "items": checked}