from __future__ import annotations

import json
import os
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple, Union

from auditwalk.core.errors import RunNotFound, RunStoreError
from auditwalk.core.paths import default_data_dir
from auditwalk.runs.schema import assert_meta_valid, assert_score_valid, normalize_meta

RUN_PREFIX = "run_"
DEFAULT_INDEX_FILE = "index.json"
DEFAULT_MAX_RUNS = 100
DEFAULT_MAX_RUN_AGE_DAYS = 30
DEFAULT_MAX_COMPARES = 200
DEFAULT_MAX_COMPARE_AGE_DAYS = 30

ARTIFACT_FILES = {
    "meta": "meta.json",
    "system": "system.json",
    "files": "files.json",
    "processes": "processes.json",
    "network": "network.json",
    "persistence": "persistence.json",
    "score": "score.json",
}


@dataclass(frozen=True)
class RunRef:
    run_id: str
    path: Path
    created_utc: int
    mode: str
    score_total: Optional[int]
    tier: Optional[str]
    label: Optional[str]
    baseline: bool


class RunStore:
    """Filesystem-backed run store with atomic writes and lightweight indexing."""

    def __init__(self, data_dir: Union[str, Path, None] = None) -> None:
        base = Path(data_dir).expanduser().resolve() if data_dir else default_data_dir("auditwalk")
        self.data_dir = base
        self.runs_dir = self.data_dir / "runs"
        self.compares_dir = self.data_dir / "compares"
        self.index_path = self.data_dir / DEFAULT_INDEX_FILE
        self.ensure_layout()

    def ensure_layout(self) -> None:
        try:
            self.runs_dir.mkdir(parents=True, exist_ok=True)
            self.compares_dir.mkdir(parents=True, exist_ok=True)
            if not self.index_path.exists():
                self._atomic_write_json(self.index_path, self._default_index())
        except PermissionError:
            # Sandbox/dev fallback: use local workspace data directory if OS path is not writable.
            fallback = Path.cwd().resolve() / ".auditwalk-data"
            self.data_dir = fallback
            self.runs_dir = self.data_dir / "runs"
            self.compares_dir = self.data_dir / "compares"
            self.index_path = self.data_dir / DEFAULT_INDEX_FILE
            self.runs_dir.mkdir(parents=True, exist_ok=True)
            self.compares_dir.mkdir(parents=True, exist_ok=True)
            if not self.index_path.exists():
                self._atomic_write_json(self.index_path, self._default_index())

    def create_run(
        self,
        run_id: str,
        artifacts: Dict[str, Any],
        *,
        mode: str,
        label: Optional[str] = None,
    ) -> Path:
        self.ensure_layout()

        if not run_id.startswith(RUN_PREFIX):
            raise RunStoreError(f"run_id must start with '{RUN_PREFIX}': {run_id}")

        run_path = self.runs_dir / run_id
        run_path.mkdir(parents=True, exist_ok=False)

        meta = artifacts.get("meta") or {}
        created_utc = int(meta.get("created_utc") or time.time())
        meta.setdefault("run_id", run_id)
        meta.setdefault("created_utc", created_utc)
        meta.setdefault("mode", mode)
        if label is not None:
            meta.setdefault("label", label)
        meta = normalize_meta(meta)
        assert_meta_valid(meta)
        self._atomic_write_json(run_path / ARTIFACT_FILES["meta"], meta)

        for key, filename in ARTIFACT_FILES.items():
            if key == "meta":
                continue
            if key in artifacts and artifacts[key] is not None:
                self._atomic_write_json(run_path / filename, artifacts[key])

        score_total, tier = self._read_score_fields(run_path)
        baseline = bool(meta.get("baseline", False))
        self._index_add(
            run_id=run_id,
            created_utc=created_utc,
            mode=mode,
            label=label,
            score_total=score_total,
            tier=tier,
            baseline=baseline,
        )
        return run_path

    def list_runs(
        self,
        *,
        last: int = 20,
        mode: str = "any",
        tier: str = "any",
        label_substr: Optional[str] = None,
    ) -> List[RunRef]:
        idx = self._load_index()
        rows = idx.get("runs", []) if isinstance(idx.get("runs", []), list) else []
        runs = [self._runref_from_index_row(r) for r in rows]

        def match(ref: RunRef) -> bool:
            if mode != "any" and ref.mode != mode:
                return False
            if tier != "any" and (ref.tier or "").lower() != tier.lower():
                return False
            if label_substr and label_substr.lower() not in (ref.label or "").lower():
                return False
            return True

        filtered = [r for r in runs if match(r)]
        filtered.sort(key=lambda r: r.created_utc, reverse=True)
        return filtered[: max(0, int(last))]

    def resolve_run(self, run_id_or_path: Union[str, Path]) -> Path:
        p = Path(run_id_or_path)
        if p.exists() and p.is_dir():
            return p.resolve()
        run_path = self.runs_dir / str(run_id_or_path)
        if run_path.exists() and run_path.is_dir():
            return run_path.resolve()
        raise RunNotFound(f"Run not found: {run_id_or_path}")

    def load_artifact(self, run_id_or_path: Union[str, Path], artifact: str) -> Any:
        if artifact not in ARTIFACT_FILES:
            raise RunStoreError(f"Unknown artifact '{artifact}'.")
        run_path = self.resolve_run(run_id_or_path)
        fpath = run_path / ARTIFACT_FILES[artifact]
        if not fpath.exists():
            raise RunNotFound(f"Artifact missing: {fpath.name}")
        return self._read_json(fpath)

    def load_run_bundle(self, run_id_or_path: Union[str, Path]) -> Dict[str, Any]:
        run_path = self.resolve_run(run_id_or_path)
        out: Dict[str, Any] = {"run_path": str(run_path)}
        for key, fname in ARTIFACT_FILES.items():
            fpath = run_path / fname
            if fpath.exists():
                out[key] = self._read_json(fpath)
        return out

    def save_score(self, run_id_or_path: Union[str, Path], score_obj: Dict[str, Any]) -> Path:
        assert_score_valid(score_obj)
        run_path = self.resolve_run(run_id_or_path)
        score_path = run_path / ARTIFACT_FILES["score"]
        self._atomic_write_json(score_path, score_obj)

        run_id = run_path.name
        idx = self._load_index()
        rows = idx.get("runs", []) if isinstance(idx.get("runs", []), list) else []
        for row in rows:
            if str(row.get("run_id", "")) == run_id:
                total = score_obj.get("total")
                tier = score_obj.get("tier")
                row["score_total"] = int(total) if isinstance(total, int) else None
                row["tier"] = str(tier) if isinstance(tier, str) else None
                break
        idx["runs"] = rows
        self._write_index(idx)
        return score_path

    def apply_retention(self, *, max_runs: int, max_age_days: int) -> Dict[str, int]:
        self.ensure_layout()
        idx = self._load_index()
        runs = idx.get("runs", []) if isinstance(idx.get("runs", []), list) else []
        runs.sort(key=lambda r: int(r.get("created_utc", 0)), reverse=True)

        now = int(time.time())
        max_age_sec = max(0, int(max_age_days)) * 86400

        keep: List[dict] = []
        deleted = 0
        for i, row in enumerate(runs):
            created = int(row.get("created_utc", 0))
            run_id = str(row.get("run_id", ""))
            too_old = (now - created) > max_age_sec if created else False
            over_count = i >= int(max_runs)
            if too_old or over_count:
                run_path = self.runs_dir / run_id
                if run_path.exists():
                    self._rm_tree(run_path)
                deleted += 1
            else:
                keep.append(row)

        idx["runs"] = keep
        self._write_index(idx)
        return {"deleted": deleted, "kept": len(keep)}

    def apply_default_retention(self) -> Dict[str, int]:
        return self.apply_retention(max_runs=DEFAULT_MAX_RUNS, max_age_days=DEFAULT_MAX_RUN_AGE_DAYS)

    def apply_compare_retention(self, *, max_compares: int, max_age_days: int) -> Dict[str, int]:
        self.ensure_layout()
        files = sorted(self.compares_dir.glob("*.json"), key=lambda p: p.stat().st_mtime, reverse=True)
        now = int(time.time())
        max_age_sec = max(0, int(max_age_days)) * 86400

        kept = 0
        deleted = 0
        for idx, path in enumerate(files):
            try:
                mtime = int(path.stat().st_mtime)
            except OSError:
                mtime = 0
            too_old = (now - mtime) > max_age_sec if mtime else False
            over_count = idx >= int(max_compares)
            if too_old or over_count:
                path.unlink(missing_ok=True)
                deleted += 1
            else:
                kept += 1
        self.sync_compare_index()
        return {"deleted": deleted, "kept": kept}

    def apply_default_compare_retention(self) -> Dict[str, int]:
        return self.apply_compare_retention(
            max_compares=DEFAULT_MAX_COMPARES,
            max_age_days=DEFAULT_MAX_COMPARE_AGE_DAYS,
        )

    def prune_orphan_compares(self) -> Dict[str, int]:
        self.ensure_layout()
        existing_runs = {p.name for p in self.runs_dir.glob("run_*") if p.is_dir()}
        deleted = 0
        kept = 0
        for diff_path in self.compares_dir.glob("diff_*.json"):
            try:
                payload = self._read_json(diff_path)
            except Exception:
                payload = {}
            left = str(payload.get("left_run_id", ""))
            right = str(payload.get("right_run_id", ""))
            if left and right and left in existing_runs and right in existing_runs:
                kept += 1
                continue
            diff_path.unlink(missing_ok=True)
            deleted += 1
        self.sync_compare_index()
        return {"deleted": deleted, "kept": kept}

    def sync_compare_index(self) -> Dict[str, int]:
        """Rebuild compare index entries from compare files on disk."""
        self.ensure_layout()
        idx = self._load_index()
        compares: List[dict[str, Any]] = []
        scanned = 0
        for diff_path in sorted(self.compares_dir.glob("diff_*.json")):
            scanned += 1
            try:
                payload = self._read_json(diff_path)
            except Exception:
                payload = {}
            compare_id = str(payload.get("compare_id") or diff_path.stem)
            left_run_id = str(payload.get("left_run_id", ""))
            right_run_id = str(payload.get("right_run_id", ""))
            generated_at = str(payload.get("generated_at", ""))
            try:
                mtime_utc = int(diff_path.stat().st_mtime)
            except OSError:
                mtime_utc = 0
            compares.append(
                {
                    "compare_id": compare_id,
                    "left_run_id": left_run_id,
                    "right_run_id": right_run_id,
                    "generated_at": generated_at,
                    "mtime_utc": mtime_utc,
                    "path": str(diff_path.resolve()),
                }
            )
        compares.sort(key=lambda row: int(row.get("mtime_utc", 0)), reverse=True)
        idx["compares"] = compares
        self._write_index(idx)
        return {"indexed": len(compares), "scanned": scanned}

    def _index_add(
        self,
        *,
        run_id: str,
        created_utc: int,
        mode: str,
        label: Optional[str],
        score_total: Optional[int],
        tier: Optional[str],
        baseline: bool,
    ) -> None:
        idx = self._load_index()
        rows = idx.get("runs", []) if isinstance(idx.get("runs", []), list) else []
        rows = [r for r in rows if str(r.get("run_id", "")) != run_id]
        rows.insert(
            0,
            {
                "run_id": run_id,
                "created_utc": int(created_utc),
                "mode": mode,
                "label": label,
                "score_total": score_total,
                "tier": tier,
                "baseline": baseline,
            },
        )
        idx["runs"] = rows
        self._write_index(idx)

    def _load_index(self) -> Dict[str, Any]:
        if not self.index_path.exists():
            return self._default_index()
        loaded = self._read_json(self.index_path)
        if not isinstance(loaded, dict):
            return self._default_index()
        out = dict(loaded)
        out["version"] = int(out.get("version", 1))
        out["runs"] = out.get("runs", []) if isinstance(out.get("runs"), list) else []
        out["compares"] = out.get("compares", []) if isinstance(out.get("compares"), list) else []
        return out

    def _write_index(self, idx: Dict[str, Any]) -> None:
        out = {
            "version": int(idx.get("version", 1)),
            "runs": idx.get("runs", []) if isinstance(idx.get("runs"), list) else [],
            "compares": idx.get("compares", []) if isinstance(idx.get("compares"), list) else [],
        }
        self._atomic_write_json(self.index_path, out)

    def _default_index(self) -> Dict[str, Any]:
        return {"version": 1, "runs": [], "compares": []}

    def _runref_from_index_row(self, row: Dict[str, Any]) -> RunRef:
        run_id = str(row.get("run_id", ""))
        return RunRef(
            run_id=run_id,
            path=self.runs_dir / run_id,
            created_utc=int(row.get("created_utc", 0)),
            mode=str(row.get("mode", "unknown")),
            score_total=row.get("score_total"),
            tier=row.get("tier"),
            label=row.get("label"),
            baseline=bool(row.get("baseline", False)),
        )

    def _read_score_fields(self, run_path: Path) -> Tuple[Optional[int], Optional[str]]:
        score_path = run_path / ARTIFACT_FILES["score"]
        if not score_path.exists():
            return None, None
        data = self._read_json(score_path)
        total = data.get("total")
        tier = data.get("tier")
        return (int(total) if total is not None else None, str(tier) if tier is not None else None)

    def _atomic_write_json(self, path: Path, data: Any) -> None:
        path.parent.mkdir(parents=True, exist_ok=True)
        tmp = path.with_suffix(path.suffix + ".tmp")
        payload = json.dumps(data, indent=2, sort_keys=True, ensure_ascii=True)
        with tmp.open("wb") as handle:
            handle.write(payload.encode("utf-8"))
            handle.flush()
            os.fsync(handle.fileno())
        os.replace(tmp, path)

    def _read_json(self, path: Path) -> Any:
        return json.loads(path.read_text(encoding="utf-8"))

    def _rm_tree(self, path: Path) -> None:
        for child in sorted(path.rglob("*"), reverse=True):
            if child.is_file() or child.is_symlink():
                child.unlink(missing_ok=True)
            elif child.is_dir():
                child.rmdir()
        path.rmdir()
