from __future__ import annotations

import hashlib
import json
import os
import re
import shutil
import tempfile
import uuid
import zipfile
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

import numpy as np

from .geometry import (
    analyse_mesh,
    export_forge_package,
    load_mesh,
    repair_mesh,
    suggest_joint_locations,
)


class ShireGateError(RuntimeError):
    pass


def utc_now() -> str:
    return datetime.now(timezone.utc).isoformat()


def slugify(value: str) -> str:
    cleaned = re.sub(r"[^A-Za-z0-9._-]+", "-", value).strip(".-_").lower()
    return cleaned[:80] or "flexiforge-project"


def sha256_file(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def canonical_hash(value: Any) -> str:
    raw = json.dumps(value, sort_keys=True, separators=(",", ":"), ensure_ascii=True).encode("utf-8")
    return hashlib.sha256(raw).hexdigest()


def _mount_points() -> list[Path]:
    points: list[Path] = []
    mounts = Path("/proc/mounts")
    if not mounts.exists():
        return points
    for line in mounts.read_text(encoding="utf-8", errors="ignore").splitlines():
        fields = line.split()
        if len(fields) >= 2:
            points.append(Path(fields[1].replace("\\040", " ")).resolve())
    return sorted(set(points), key=lambda p: len(str(p)), reverse=True)


def verify_network_vault(root: Path) -> Path:
    resolved = root.expanduser().resolve()
    if not resolved.exists() or not resolved.is_dir():
        raise ShireGateError(f"SHiREVault is unavailable: {resolved}")

    allow_unmounted = os.environ.get("FLEXIFORGE_ALLOW_UNMOUNTED_VAULT") == "1"
    if not allow_unmounted:
        mounted_under = next((m for m in _mount_points() if resolved == m or resolved.is_relative_to(m)), None)
        if mounted_under is None or mounted_under == Path("/"):
            raise ShireGateError(
                f"SHiREVault path is not backed by a verified non-root mount: {resolved}. "
                "FlexiForge stopped rather than exporting locally."
            )

    try:
        probe_dir = resolved / ".flexiforge-write-test"
        probe_dir.mkdir(parents=True, exist_ok=True)
        with tempfile.NamedTemporaryFile(prefix="probe-", dir=probe_dir, delete=True) as probe:
            probe.write(b"FlexiForge SHiREVault verification")
            probe.flush()
        try:
            probe_dir.rmdir()
        except OSError:
            pass
    except Exception as exc:
        raise ShireGateError(f"SHiREVault is not writable: {resolved}: {exc}") from exc
    return resolved


class ShireJobStore:
    STATES = {
        "REQUESTED",
        "AWAITING_DESIGN_SELECTION",
        "AWAITING_DESIGN_APPROVAL",
        "APPROVED_FOR_FORGE",
        "FORGED",
        "VALIDATED",
        "AWAITING_EXPORT_APPROVAL",
        "APPROVED_FOR_EXPORT",
        "EXPORTED",
        "FAILED",
    }

    def __init__(self, jobs_root: Path, models_root: Path, exports_root: Path):
        self.jobs_root = jobs_root
        self.models_root = models_root
        self.exports_root = exports_root
        for path in (jobs_root, models_root, exports_root):
            path.mkdir(parents=True, exist_ok=True)

    def _job_dir(self, job_id: str) -> Path:
        return self.jobs_root / job_id

    def _manifest_path(self, job_id: str) -> Path:
        return self._job_dir(job_id) / "manifest.json"

    def _save(self, job: dict[str, Any]) -> dict[str, Any]:
        job["updated_at"] = utc_now()
        job["manifest_hash"] = canonical_hash({k: v for k, v in job.items() if k != "manifest_hash"})
        path = self._manifest_path(job["job_id"])
        path.parent.mkdir(parents=True, exist_ok=True)
        tmp = path.with_suffix(".tmp")
        tmp.write_text(json.dumps(job, indent=2, sort_keys=True), encoding="utf-8")
        tmp.replace(path)
        return job

    def get(self, job_id: str) -> dict[str, Any]:
        path = self._manifest_path(job_id)
        if not path.exists():
            raise FileNotFoundError(f"SHiRE FlexiForge job not found: {job_id}")
        return json.loads(path.read_text(encoding="utf-8"))

    def _event(self, job: dict[str, Any], action: str, detail: str) -> None:
        job.setdefault("events", []).append({"time": utc_now(), "action": action, "detail": detail})

    def _require(self, job: dict[str, Any], *states: str) -> None:
        if job["state"] not in states:
            raise ShireGateError(f"Action blocked in state {job['state']}; required: {', '.join(states)}")

    def create_job(
        self,
        model_id: str,
        project_name: str,
        intent: str,
        movement: str,
        clearance_mm: float,
        max_suggestions: int,
    ) -> dict[str, Any]:
        source = self.models_root / f"{model_id}.stl"
        if not source.exists():
            raise FileNotFoundError("Uploaded model not found.")
        mesh = load_mesh(source)
        mesh, notes = repair_mesh(mesh)
        analysis = analyse_mesh(mesh, notes).__dict__
        suggestions = suggest_joint_locations(mesh, movement, clearance_mm, max_suggestions)
        if not suggestions:
            raise ShireGateError("SHiRE could not generate a usable automatic joint location.")

        meta_path = self.models_root / f"{model_id}.json"
        original_name = source.name
        if meta_path.exists():
            original_name = json.loads(meta_path.read_text(encoding="utf-8")).get("original_name", original_name)

        job_id = f"FF-{datetime.now().strftime('%Y%m%d')}-{uuid.uuid4().hex[:8].upper()}"
        job = {
            "schema_version": 1,
            "engine": "FlexiForge SHiRE Bridge 0.2.0",
            "job_id": job_id,
            "project_name": project_name,
            "project_slug": slugify(project_name),
            "intent": intent,
            "state": "AWAITING_DESIGN_SELECTION",
            "created_at": utc_now(),
            "updated_at": utc_now(),
            "model": {
                "model_id": model_id,
                "original_name": original_name,
                "path": str(source),
                "sha256": sha256_file(source),
                "analysis": analysis,
            },
            "request": {"movement": movement, "clearance_mm": clearance_mm},
            "suggestions": suggestions,
            "selected_joint": None,
            "approvals": {"design": None, "export": None},
            "forge": None,
            "validation": None,
            "export": None,
            "events": [],
        }
        self._event(job, "CREATE_JOB", "SHiRE analysed the STL and generated deterministic joint suggestions.")
        return self._save(job)

    def select_suggestion(self, job_id: str, suggestion_id: str, requested_type: str = "auto") -> dict[str, Any]:
        job = self.get(job_id)
        self._require(job, "AWAITING_DESIGN_SELECTION", "AWAITING_DESIGN_APPROVAL")
        suggestion = next((s for s in job["suggestions"] if s["suggestion_id"] == suggestion_id), None)
        if suggestion is None:
            raise ShireGateError(f"Unknown suggestion: {suggestion_id}")
        joint_type = suggestion["recommended_joint"]["type"] if requested_type == "auto" else requested_type
        job["selected_joint"] = {
            "suggestion_id": suggestion_id,
            "point": suggestion["point"],
            "cut_normal": suggestion["cut_normal"],
            "joint_type": joint_type,
            "clearance_mm": job["request"]["clearance_mm"],
            "confidence": suggestion["confidence"],
        }
        job["approvals"]["design"] = None
        job["state"] = "AWAITING_DESIGN_APPROVAL"
        self._event(job, "SELECT_DESIGN", f"Selected {suggestion_id} using {joint_type}; design approval required.")
        return self._save(job)

    def _design_basis(self, job: dict[str, Any]) -> dict[str, Any]:
        return {
            "job_id": job["job_id"],
            "model_sha256": job["model"]["sha256"],
            "request": job["request"],
            "selected_joint": job["selected_joint"],
        }

    def _export_basis(self, job: dict[str, Any]) -> dict[str, Any]:
        return {
            "job_id": job["job_id"],
            "design_approval": job["approvals"]["design"],
            "forge": job["forge"],
            "validation": job["validation"],
        }

    def approval_token(self, job_id: str, kind: str) -> str:
        job = self.get(job_id)
        if kind == "design":
            self._require(job, "AWAITING_DESIGN_APPROVAL")
            basis = self._design_basis(job)
            prefix = "APPROVE-DESIGN"
        elif kind == "export":
            self._require(job, "VALIDATED", "AWAITING_EXPORT_APPROVAL")
            basis = self._export_basis(job)
            prefix = "APPROVE-EXPORT"
            if job["state"] == "VALIDATED":
                job["state"] = "AWAITING_EXPORT_APPROVAL"
                self._event(job, "REQUEST_EXPORT_APPROVAL", "Validated output is awaiting Ray's export approval.")
                self._save(job)
        else:
            raise ValueError("Approval kind must be design or export.")
        fingerprint = canonical_hash(basis)[:10].upper()
        return f"{prefix}-{job['job_id'][-8:]}-{fingerprint}"

    def approve(self, job_id: str, kind: str, token: str) -> dict[str, Any]:
        expected = self.approval_token(job_id, kind)
        if token.strip() != expected:
            raise ShireGateError("Approval token does not match the current immutable job design.")
        job = self.get(job_id)
        if kind == "design":
            job["approvals"]["design"] = {
                "token": expected,
                "approved_at": utc_now(),
                "basis_hash": canonical_hash(self._design_basis(job)),
            }
            job["state"] = "APPROVED_FOR_FORGE"
            self._event(job, "APPROVE_DESIGN", "Ray approved the selected joint design for geometry generation.")
        else:
            job["approvals"]["export"] = {
                "token": expected,
                "approved_at": utc_now(),
                "basis_hash": canonical_hash(self._export_basis(job)),
            }
            job["state"] = "APPROVED_FOR_EXPORT"
            self._event(job, "APPROVE_EXPORT", "Ray approved export to the SHiREVault network drive.")
        return self._save(job)

    def forge(self, job_id: str) -> dict[str, Any]:
        job = self.get(job_id)
        self._require(job, "APPROVED_FOR_FORGE")
        basis_hash = canonical_hash(self._design_basis(job))
        if job["approvals"]["design"]["basis_hash"] != basis_hash:
            raise ShireGateError("Design changed after approval; forge approval is invalid.")

        mesh = load_mesh(Path(job["model"]["path"]))
        mesh, _ = repair_mesh(mesh)
        selected = job["selected_joint"]
        point = np.array([selected["point"][k] for k in ("x", "y", "z")], dtype=float)
        normal = np.array([selected["cut_normal"][k] for k in ("x", "y", "z")], dtype=float)
        output_root = self._job_dir(job_id) / "forge_output"
        output_root.mkdir(parents=True, exist_ok=True)
        result = export_forge_package(
            output_dir=output_root,
            source_name=job["project_slug"],
            mesh=mesh,
            point=point,
            normal=normal,
            joint_type=selected["joint_type"],
            clearance=selected["clearance_mm"],
            layout_gap=10.0,
        )
        package_dir: Path = result.pop("directory")
        zip_path = output_root / f"{result['package_id']}.zip"
        with zipfile.ZipFile(zip_path, "w", compression=zipfile.ZIP_DEFLATED) as archive:
            for file_path in sorted(package_dir.iterdir()):
                archive.write(file_path, arcname=file_path.name)

        job["forge"] = {
            "package_id": result["package_id"],
            "directory": str(package_dir),
            "zip_path": str(zip_path),
            "files": result["files"],
            "report": result["report"],
            "zip_sha256": sha256_file(zip_path),
            "forged_at": utc_now(),
        }
        job["state"] = "FORGED"
        self._event(job, "FORGE", "FlexiForge generated the approved articulation package locally; export remains blocked.")
        return self._save(job)

    def validate(self, job_id: str) -> dict[str, Any]:
        job = self.get(job_id)
        self._require(job, "FORGED")
        package_dir = Path(job["forge"]["directory"])
        checks: list[dict[str, Any]] = []
        warnings: list[str] = []
        for key, filename in job["forge"]["files"].items():
            path = package_dir / filename
            if not path.exists():
                raise ShireGateError(f"Validation failed: missing {filename}")
            mesh = load_mesh(path)
            item = {
                "role": key,
                "file": filename,
                "sha256": sha256_file(path),
                "vertices": int(len(mesh.vertices)),
                "faces": int(len(mesh.faces)),
                "watertight": bool(mesh.is_watertight),
                "extents_mm": [round(float(v), 3) for v in mesh.extents],
            }
            if len(mesh.faces) == 0 or float(max(mesh.extents)) <= 0:
                raise ShireGateError(f"Validation failed: empty or zero-sized output {filename}")
            if not mesh.is_watertight:
                warnings.append(f"{filename} is not watertight and needs slicer inspection.")
            checks.append(item)

        job["validation"] = {
            "status": "PASS_WITH_WARNINGS" if warnings else "PASS",
            "validated_at": utc_now(),
            "checks": checks,
            "warnings": warnings,
            "required_human_checks": [
                "Inspect every part in the slicer.",
                "Confirm the model is oriented correctly and remains on the build plate.",
                "Print a clearance test before a long production run.",
            ],
        }
        job["state"] = "VALIDATED"
        self._event(job, "VALIDATE", f"Geometry validation completed: {job['validation']['status']}.")
        return self._save(job)

    def export_to_vault(self, job_id: str, vault_root: Path | None = None) -> dict[str, Any]:
        job = self.get(job_id)
        self._require(job, "APPROVED_FOR_EXPORT")
        basis_hash = canonical_hash(self._export_basis(job))
        if job["approvals"]["export"]["basis_hash"] != basis_hash:
            raise ShireGateError("Validated output changed after approval; export approval is invalid.")

        configured = vault_root or Path(os.environ.get("SHIREVAULT_ROOT", "/SHiREVault"))
        root = verify_network_vault(configured)
        destination = root / "SHIRE3D" / "FlexiForge" / "Projects" / job["project_slug"] / job_id
        if destination.exists():
            raise ShireGateError(f"Export destination already exists: {destination}")
        destination.mkdir(parents=True, exist_ok=False)

        package_dir = Path(job["forge"]["directory"])
        exported_package = destination / "print_package"
        shutil.copytree(package_dir, exported_package)
        shutil.copy2(job["forge"]["zip_path"], destination / Path(job["forge"]["zip_path"]).name)

        final_manifest = dict(job)
        final_manifest["state"] = "EXPORTED"
        final_manifest["export"] = {
            "destination": str(destination),
            "vault_root": str(root),
            "exported_at": utc_now(),
            "verified_files": {},
        }
        for path in sorted(destination.rglob("*")):
            if path.is_file():
                final_manifest["export"]["verified_files"][str(path.relative_to(destination))] = sha256_file(path)
        (destination / "SHIRE_FlexiForge_Manifest.json").write_text(
            json.dumps(final_manifest, indent=2, sort_keys=True), encoding="utf-8"
        )
        # Verify the manifest itself after writing it.
        final_manifest["export"]["verified_files"]["SHIRE_FlexiForge_Manifest.json"] = sha256_file(
            destination / "SHIRE_FlexiForge_Manifest.json"
        )

        job["state"] = "EXPORTED"
        job["export"] = final_manifest["export"]
        self._event(job, "EXPORT", f"Verified export completed to network drive: {destination}")
        return self._save(job)
