"""Snapshot service: wrap SurfaceDriver.dump into deduped UIState assets."""

from __future__ import annotations

import time
from dataclasses import asdict, is_dataclass
from pathlib import Path
from typing import Any, Optional

from models.observation import UIObservation
from models.serialize import to_dict, write_model_json
from models.state import UIState, compute_state_fingerprint
from redaction import (
    redact_locatable_rows,
    sanitize_artifact,
    screenshot_sensitivity_meta,
)


def _dump_to_parts(dump: Any) -> dict[str, Any]:
    if dump is None:
        return {}
    if isinstance(dump, dict):
        return dict(dump)
    # DumpResult dataclass or duck-typed dump objects (tests / drivers).
    if is_dataclass(dump) and not isinstance(dump, type) or hasattr(dump, "surface"):
        return {
            "surface": getattr(dump, "surface", "") or "",
            "driver": getattr(dump, "driver", "") or "",
            "window": dict(getattr(dump, "window", None) or {}),
            "tree": getattr(dump, "tree", None) or {},
            "index": list(getattr(dump, "index", None) or []),
            "locatable": list(getattr(dump, "locatable", None) or []),
            "stats": dict(getattr(dump, "stats", None) or {}),
            "selection": dict(getattr(dump, "selection", None) or {}),
            "url": getattr(dump, "url", "") or "",
            "page_title": getattr(dump, "page_title", "") or "",
            "screenshot_bytes": getattr(dump, "screenshot_bytes", None),
            "use_hwnd_screenshot": bool(getattr(dump, "use_hwnd_screenshot", False)),
            "locatable_text": getattr(dump, "locatable_text", "") or "",
            "tree_text": getattr(dump, "tree_text", "") or "",
        }
    return {}


class SnapshotService:
    """
    Creates UIState records with fingerprint dedupe.

    Reuses SurfaceDriver.dump(); does not reimplement browser/desktop dump.
    """

    def __init__(
        self,
        flow_dir: Path,
        *,
        save_sensitive_screenshots: bool = False,
        persist_evidence: bool = True,
    ) -> None:
        self.flow_dir = Path(flow_dir)
        self.states_dir = self.flow_dir / "states"
        self.states_dir.mkdir(parents=True, exist_ok=True)
        self.observations_dir = self.flow_dir / "observations"
        self.observations_dir.mkdir(parents=True, exist_ok=True)
        self.save_sensitive_screenshots = save_sensitive_screenshots
        self.persist_evidence = persist_evidence
        self._by_fp: dict[str, UIState] = {}
        self._order: list[str] = []
        self.latest: Optional[UIState] = None
        self.latest_observation: Optional[UIObservation] = None
        self._observation_order: list[str] = []

    def load_existing(self) -> None:
        """Index existing state JSON files (best-effort)."""
        if not self.states_dir.exists():
            return
        for path in sorted(self.states_dir.glob("*.json")):
            try:
                import json

                data = json.loads(path.read_text(encoding="utf-8"))
                st = UIState(
                    id=data.get("id") or path.stem,
                    surface=data.get("surface") or "windows",
                    driver=data.get("driver") or "",
                    captured_at=data.get("captured_at") or "",
                    fingerprint=data.get("fingerprint") or "",
                    previous_state_id=data.get("previous_state_id"),
                    change_reason=data.get("change_reason") or "",
                    app=data.get("app") or {},
                    window=data.get("window") or {},
                    url=data.get("url") or "",
                    page_title=data.get("page_title") or "",
                    tree_path=data.get("tree_path") or "",
                    index_path=data.get("index_path") or "",
                    locatable_path=data.get("locatable_path") or "",
                    screenshot_path=data.get("screenshot_path") or "",
                    stats=data.get("stats") or {},
                    selection=data.get("selection") or {},
                    duplicate_of=data.get("duplicate_of"),
                    observation_count=int(data.get("observation_count") or 1),
                    may_contain_sensitive_data=bool(
                        data.get("may_contain_sensitive_data")
                    ),
                )
                if st.fingerprint:
                    self._by_fp[st.fingerprint] = st
                    if st.id not in self._order:
                        self._order.append(st.id)
                    self.latest = st
            except Exception:
                continue
        for path in sorted(self.observations_dir.glob("*/meta.json")):
            try:
                import json

                data = json.loads(path.read_text(encoding="utf-8"))
                obs = UIObservation(
                    **{k: v for k, v in data.items() if k in UIObservation.__dataclass_fields__}
                )
                self._observation_order.append(obs.id)
                self.latest_observation = obs
            except Exception:
                continue

    def capture_from_dump(
        self,
        dump: Any,
        *,
        change_reason: str = "",
        force_persist: bool = False,
        previous_state_id: Optional[str] = None,
        redact: bool = True,
    ) -> UIState:
        parts = _dump_to_parts(dump)
        locatable = list(parts.get("locatable") or [])
        sensitive = False
        if redact:
            locatable, sensitive = redact_locatable_rows(locatable)

        window = dict(parts.get("window") or {})
        url = parts.get("url") or ""
        page_title = parts.get("page_title") or ""
        surface = parts.get("surface") or "windows"
        driver = parts.get("driver") or ""
        fp = compute_state_fingerprint(
            surface=surface,
            driver=driver,
            url=url,
            page_title=page_title,
            window=window,
            locatable=locatable,
            app={
                "title": window.get("title"),
                "class_name": window.get("class_name"),
                "process_name": window.get("process_name"),
            },
        )

        existing = self._by_fp.get(fp)
        if existing is not None:
            existing.bump_observation()
            existing.duplicate_of = None
            # Update meta path only — do not rewrite full tree evidence.
            write_model_json(self.states_dir / f"{existing.id}.json", existing)
            self.latest = existing
            self._persist_observation(
                existing,
                parts,
                locatable=locatable,
                sensitive=sensitive,
                change_reason=change_reason or "capture",
            )
            return existing

        prev = previous_state_id
        if prev is None and self.latest is not None:
            prev = self.latest.id

        state = UIState(
            surface=surface,
            driver=driver,
            fingerprint=fp,
            previous_state_id=prev,
            change_reason=change_reason or "capture",
            app={
                "title": window.get("title"),
                "class_name": window.get("class_name"),
                "process_name": window.get("process_name"),
                "pid": window.get("pid"),
                "hwnd": window.get("hwnd"),
            },
            window=window,
            url=url,
            page_title=page_title,
            stats=dict(parts.get("stats") or {}),
            selection=dict(parts.get("selection") or {}),
            locatable=locatable,
            may_contain_sensitive_data=sensitive,
        )
        state.ensure_fingerprint()

        should_write_evidence = force_persist or bool(change_reason) or self.persist_evidence
        if should_write_evidence:
            self._persist_evidence(state, parts, sensitive=sensitive)

        write_model_json(self.states_dir / f"{state.id}.json", state)
        self._by_fp[fp] = state
        self._order.append(state.id)
        self.latest = state
        self._persist_observation(
            state,
            parts,
            locatable=locatable,
            sensitive=sensitive,
            change_reason=change_reason or "capture",
        )
        return state

    def _persist_observation(
        self,
        state: UIState,
        parts: dict[str, Any],
        *,
        locatable: list[dict[str, Any]],
        sensitive: bool,
        change_reason: str,
    ) -> UIObservation:
        stats = dict(parts.get("stats") or {})
        selection = dict(parts.get("selection") or {})
        active_page = dict(stats.get("active_page") or {})
        browser_session = dict(stats.get("browser_session") or {})
        obs = UIObservation(
            state_id=state.id,
            previous_observation_id=(
                self.latest_observation.id if self.latest_observation else None
            ),
            change_reason=change_reason,
            surface=state.surface,
            driver=state.driver,
            page_id=str(
                active_page.get("page_id")
                or browser_session.get("page_id")
                or selection.get("page_id")
                or ""
            ),
            document_id=str(selection.get("document_id") or ""),
            frame_id=str(selection.get("frame_id") or ""),
            url=state.url,
            page_title=state.page_title,
            stats=stats,
            selection=selection,
            may_contain_sensitive_data=sensitive,
        )
        evidence_dir = self.observations_dir / obs.id
        evidence_dir.mkdir(parents=True, exist_ok=True)
        import artifacts as art

        art.write_json(
            evidence_dir / "tree.json",
            sanitize_artifact(parts.get("tree") or {}),
        )
        clean_index = sanitize_artifact(list(parts.get("index") or []))
        art.write_jsonl(evidence_dir / "index.jsonl", list(clean_index or []))
        art.write_json(evidence_dir / "locatable.json", list(locatable or []))
        obs.tree_path = f"observations/{obs.id}/tree.json"
        obs.index_path = f"observations/{obs.id}/index.jsonl"
        obs.locatable_path = f"observations/{obs.id}/locatable.json"

        sensitivity_meta = screenshot_sensitivity_meta(
            may_contain_sensitive=sensitive,
            save_sensitive_screenshots=self.save_sensitive_screenshots,
        )
        art.write_json(evidence_dir / "sensitivity.json", sensitivity_meta)
        if (not sensitive) or self.save_sensitive_screenshots:
            png = parts.get("screenshot_bytes")
            if png:
                try:
                    (evidence_dir / "screen.png").write_bytes(png)
                    obs.screenshot_path = f"observations/{obs.id}/screen.png"
                except Exception:
                    pass
            elif parts.get("use_hwnd_screenshot"):
                hwnd = int((parts.get("window") or {}).get("hwnd") or 0)
                if hwnd and art.capture_window_png(hwnd, evidence_dir / "screen.png"):
                    obs.screenshot_path = f"observations/{obs.id}/screen.png"

        write_model_json(evidence_dir / "meta.json", obs)
        self.latest_observation = obs
        self._observation_order.append(obs.id)
        return obs

    def _persist_evidence(
        self,
        state: UIState,
        parts: dict[str, Any],
        *,
        sensitive: bool,
    ) -> None:
        evidence_dir = self.states_dir / state.id
        evidence_dir.mkdir(parents=True, exist_ok=True)
        import artifacts as art

        art.write_json(
            evidence_dir / "tree.json",
            sanitize_artifact(parts.get("tree") or {}),
        )
        clean_index = sanitize_artifact(list(parts.get("index") or []))
        art.write_jsonl(evidence_dir / "index.jsonl", list(clean_index or []))
        art.write_json(evidence_dir / "locatable.json", list(state.locatable or []))
        state.tree_path = f"states/{state.id}/tree.json"
        state.index_path = f"states/{state.id}/index.jsonl"
        state.locatable_path = f"states/{state.id}/locatable.json"

        sens_meta = screenshot_sensitivity_meta(
            may_contain_sensitive=sensitive,
            save_sensitive_screenshots=self.save_sensitive_screenshots,
        )
        art.write_json(evidence_dir / "sensitivity.json", sens_meta)

        save_png = (not sensitive) or self.save_sensitive_screenshots
        if save_png:
            png = parts.get("screenshot_bytes")
            if png:
                try:
                    (evidence_dir / "screen.png").write_bytes(png)
                    state.screenshot_path = f"states/{state.id}/screen.png"
                except Exception:
                    pass
            elif parts.get("use_hwnd_screenshot"):
                hwnd = int((parts.get("window") or {}).get("hwnd") or 0)
                if hwnd:
                    ok = art.capture_window_png(hwnd, evidence_dir / "screen.png")
                    if ok:
                        state.screenshot_path = f"states/{state.id}/screen.png"
        elif sensitive:
            state.may_contain_sensitive_data = True

    def wait_until_stable(
        self,
        dump_fn,
        *,
        timeout_sec: float = 3.0,
        poll_sec: float = 0.25,
        change_reason: str = "after_action",
    ) -> UIState:
        """
        Poll dump_fn until two consecutive fingerprints match or timeout.
        dump_fn() -> DumpResult / dump-like object
        """
        deadline = time.monotonic() + max(0.2, float(timeout_sec))
        last_fp: Optional[str] = None
        last_dump = None
        while True:
            last_dump = dump_fn()
            parts = _dump_to_parts(last_dump)
            locatable, _ = redact_locatable_rows(list(parts.get("locatable") or []))
            fp = compute_state_fingerprint(
                surface=parts.get("surface") or "windows",
                driver=parts.get("driver") or "",
                url=parts.get("url") or "",
                page_title=parts.get("page_title") or "",
                window=parts.get("window") or {},
                locatable=locatable,
            )
            if last_fp is not None and fp == last_fp:
                break
            last_fp = fp
            if time.monotonic() >= deadline:
                break
            time.sleep(max(0.05, float(poll_sec)))
        return self.capture_from_dump(
            last_dump,
            change_reason=change_reason,
            force_persist=True,
        )

    def state_ids(self) -> list[str]:
        return list(self._order)

    def observation_ids(self) -> list[str]:
        return list(self._observation_order)

    def get(self, state_id: str) -> Optional[UIState]:
        for st in self._by_fp.values():
            if st.id == state_id:
                return st
        path = self.states_dir / f"{state_id}.json"
        if path.exists():
            import json
            from models.state import UIState as S

            data = json.loads(path.read_text(encoding="utf-8"))
            return S(**{k: v for k, v in data.items() if k in S.__dataclass_fields__})
        return None

    def to_summary_dict(self, state: UIState) -> dict[str, Any]:
        d = to_dict(state)
        d["locatable"] = list(state.locatable or [])
        return d