from __future__ import annotations

import argparse
import json
from pathlib import Path
from typing import Any


ROOT = Path(__file__).resolve().parents[1]
DEFAULT_MANIFEST = ROOT / "reports" / "llm-artifact-manifest.json"
DEFAULT_OUTPUT = ROOT / "reports" / "llm-artifact-check.json"


def load_json(path: Path) -> dict[str, Any]:
    data = json.loads(path.read_text(encoding="utf-8-sig"))
    if not isinstance(data, dict):
        raise ValueError(f"{path} must contain a JSON object")
    return data


def iter_files(path: Path) -> list[Path]:
    if not path.exists():
        return []
    return sorted(item for item in path.rglob("*") if item.is_file())


def resolve_record_path(record_path: str, *, source_root: Path, target_root: Path | None) -> Path:
    path = Path(record_path)
    if target_root is None:
        return path
    try:
        relative = path.relative_to(source_root)
        return target_root / relative
    except ValueError:
        return target_root / path.name


def check_artifact(record: dict[str, Any], *, source_root: Path, target_root: Path | None, strict_counts: bool) -> list[dict[str, Any]]:
    findings: list[dict[str, Any]] = []
    path = resolve_record_path(str(record.get("path") or ""), source_root=source_root, target_root=target_root)
    expected_exists = bool(record.get("exists"))
    if expected_exists and not path.exists():
        return [{"severity": "error", "code": "artifact_missing", "artifact": record.get("name"), "path": str(path)}]
    if not path.exists():
        return []
    if not path.is_dir():
        findings.append({"severity": "error", "code": "artifact_not_directory", "artifact": record.get("name"), "path": str(path)})
        return findings

    files = iter_files(path)
    size = sum(item.stat().st_size for item in files)
    expected_count = int(record.get("file_count") or 0)
    expected_size = int(record.get("total_size_bytes") or 0)
    if strict_counts and len(files) != expected_count:
        findings.append(
            {
                "severity": "error",
                "code": "file_count_mismatch",
                "artifact": record.get("name"),
                "path": str(path),
                "expected": expected_count,
                "actual": len(files),
            }
        )
    elif len(files) < expected_count:
        findings.append(
            {
                "severity": "warning",
                "code": "file_count_decreased",
                "artifact": record.get("name"),
                "path": str(path),
                "expected_at_least": expected_count,
                "actual": len(files),
            }
        )
    if strict_counts and size != expected_size:
        findings.append(
            {
                "severity": "error",
                "code": "total_size_mismatch",
                "artifact": record.get("name"),
                "path": str(path),
                "expected": expected_size,
                "actual": size,
            }
        )
    elif size < expected_size:
        findings.append(
            {
                "severity": "warning",
                "code": "total_size_decreased",
                "artifact": record.get("name"),
                "path": str(path),
                "expected_at_least": expected_size,
                "actual": size,
            }
        )
    return findings


def check_manifest(manifest: dict[str, Any], *, target_root: Path | None, strict_counts: bool) -> dict[str, Any]:
    source_root = Path(str(manifest.get("workspace_root") or ROOT))
    findings: list[dict[str, Any]] = []
    for record in manifest.get("artifacts") or []:
        findings.extend(check_artifact(record, source_root=source_root, target_root=target_root, strict_counts=strict_counts))

    errors = [item for item in findings if item.get("severity") == "error"]
    warnings = [item for item in findings if item.get("severity") == "warning"]
    return {
        "schema": "llm_artifact_manifest_check.v1",
        "manifest_schema": manifest.get("schema"),
        "manifest_created_at": manifest.get("created_at"),
        "target_root": str(target_root) if target_root else None,
        "strict_counts": strict_counts,
        "passed": not errors,
        "counts": {
            "artifacts": len(manifest.get("artifacts") or []),
            "errors": len(errors),
            "warnings": len(warnings),
            "findings": len(findings),
        },
        "findings": findings,
    }


def main() -> int:
    parser = argparse.ArgumentParser(description="Check that local LLM/RAG artifacts from a manifest are present after transfer or before Docker launch.")
    parser.add_argument("--manifest", type=Path, default=DEFAULT_MANIFEST)
    parser.add_argument("--output", type=Path, default=DEFAULT_OUTPUT)
    parser.add_argument("--target-root", type=Path, help="New workspace root after transfer. If omitted, paths are checked as recorded.")
    parser.add_argument("--strict-counts", action="store_true", help="Fail when file counts or total sizes differ exactly.")
    args = parser.parse_args()

    result = check_manifest(load_json(args.manifest), target_root=args.target_root, strict_counts=args.strict_counts)
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps(result, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
    print(json.dumps({"output": str(args.output), "passed": result["passed"], "counts": result["counts"]}, ensure_ascii=False))
    return 0 if result["passed"] else 2


if __name__ == "__main__":
    raise SystemExit(main())