"""Load precomputed demo traces for Streamlit replay mode."""

from __future__ import annotations

import json
import re
from pathlib import Path
from typing import Any, Optional

REPO_ROOT = Path(__file__).resolve().parent.parent
REPLAY_ROOT = REPO_ROOT / "replay_traces"


def _load_json(path: Path) -> Any | None:
    if not path.exists():
        return None
    try:
        return json.loads(path.read_text(encoding="utf-8"))
    except Exception:
        return None


def parse_arxiv_id(paper_input: str) -> str | None:
    s = (paper_input or "").strip()
    if not s:
        return None
    if "arxiv.org" in s:
        m = re.search(r"arxiv\.org/(abs|pdf)/([^/?#]+)", s)
        if not m:
            return None
        s = m.group(2)
    s = s.replace(".pdf", "")
    s = re.sub(r"v\d+$", "", s)
    if not re.match(r"^[0-9]{4}\.[0-9]{4,5}$", s):
        return None
    return s


def load_index() -> dict:
    data = _load_json(REPLAY_ROOT / "index.json")
    return data if isinstance(data, dict) else {"examples": []}


def list_replay_examples() -> list[dict]:
    examples = load_index().get("examples") or []
    out = []
    for item in examples:
        if not isinstance(item, dict):
            continue
        arxiv_id = item.get("arxiv_id")
        if not arxiv_id:
            continue
        meta_path = REPLAY_ROOT / str(arxiv_id) / "replay_meta.json"
        if not meta_path.exists():
            continue
        out.append(item)
    return out


def find_trace_dir(paper_input: str) -> Optional[Path]:
    arxiv_id = parse_arxiv_id(paper_input)
    if not arxiv_id:
        return None
    trace_dir = REPLAY_ROOT / arxiv_id
    if (trace_dir / "replay_meta.json").exists():
        return trace_dir
    return None


def load_trace(paper_input: str) -> Optional[dict]:
    """Return a session-ready payload for a saved demo trace, or None."""
    trace_dir = find_trace_dir(paper_input)
    if trace_dir is None:
        return None

    meta = _load_json(trace_dir / "replay_meta.json")
    if not isinstance(meta, dict):
        return None

    paper_rel = meta.get("paper_dir") or f"processed_papers/{meta.get('arxiv_id')}"
    paper_dir = (trace_dir / paper_rel).resolve()
    if not paper_dir.exists():
        return None

    payload_rel = meta.get("annotation_payload_path")
    payload_path = (trace_dir / payload_rel).resolve() if payload_rel else None
    if payload_path and not payload_path.exists():
        payload_path = None

    events = []
    for line in meta.get("events") or []:
        text = str(line).strip()
        if not text:
            continue
        # Normalize absolute annotation paths for display.
        if text.startswith("[annotation] complete:"):
            events.append("Step 8 complete")
            events.append("Pipeline completed successfully.")
            continue
        if text == "Pipeline completed successfully.":
            # Keep a single completion marker after annotation normalization.
            if events and events[-1] == "Pipeline completed successfully.":
                continue
        events.append(text)

    # Ensure a readable end marker for completed traces.
    status = str(meta.get("status") or "Completed")
    if status == "Completed" and "Pipeline completed successfully." not in events:
        events.append("Pipeline completed successfully.")

    return {
        "label": meta.get("label") or meta.get("arxiv_id"),
        "arxiv_id": meta.get("arxiv_id"),
        "paper_input": meta.get("paper_input") or f"https://arxiv.org/abs/{meta.get('arxiv_id')}",
        "status": status,
        "run_status": status,
        "run_events": events,
        "run_logs": list(events),
        "run_dir_path": str(trace_dir),
        "paper_dir_path": str(paper_dir),
        "annotation_payload_path": str(payload_path) if payload_path else None,
        "annotation_skipped_reason": meta.get("annotation_skipped_reason"),
        "pipeline_failed_reason": meta.get("pipeline_failed_reason"),
        "pipeline_stopped_reason": meta.get("pipeline_stopped_reason"),
        "artifact_path": str(trace_dir / "summary.txt") if (trace_dir / "summary.txt").exists() else None,
        "remote_artifact_ref": "",
        "replay_mode": True,
        "cluster_count": meta.get("cluster_count"),
        "claim_count": meta.get("claim_count"),
    }