#!/usr/bin/env python3
"""Fast relevance screening for a broad Agent-memory research corpus."""

from __future__ import annotations

import argparse
import json
import re
import time
import urllib.error
import urllib.request
from datetime import date
from pathlib import Path


ROOT = Path(__file__).resolve().parents[2]

PROMPT = """Screen papers for a research review about memory in LLM-based agents.
Use title and abstract only. Return JSON with a `papers` array in input order.
For each paper return:
- arxiv_id: exact input id
- relevance: core if memory is the main research object; supporting if memory is a substantial
  mechanism, evaluation dimension, or failure mode; peripheral if memory is incidental, means
  GPU/KV memory only, refers to human memory only, or the paper is mainly an unrelated application
- problem: choose one: access, representation, update_forgetting, execution_state, experience_learning,
  evaluation, security_privacy, shared_memory, systems_cost, theory, application, unrelated
- reason: at most 12 words
Be conservative and output JSON only."""

RELEVANCE_VALUES = {"core", "supporting", "peripheral"}
PROBLEM_VALUES = {
    "access",
    "representation",
    "update_forgetting",
    "execution_state",
    "experience_learning",
    "evaluation",
    "security_privacy",
    "shared_memory",
    "systems_cost",
    "theory",
    "application",
    "unrelated",
}
PROBLEM_ALIASES = {
    "adaptation": "experience_learning",
    "continual_learning": "experience_learning",
    "episodic": "representation",
    "episodic_memory": "representation",
    "experience": "experience_learning",
    "forgetting": "update_forgetting",
    "implicit_personalization": "application",
    "memory": "representation",
    "privacy": "security_privacy",
    "security": "security_privacy",
}


def write_json(path: Path, payload: dict) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    tmp = path.with_suffix(path.suffix + ".tmp")
    tmp.write_text(json.dumps(payload, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
    tmp.replace(path)


def generate(url: str, model: str, records: list[dict], timeout: int, retries: int) -> list[dict]:
    compact = [
        {"arxiv_id": item["arxiv_id"], "title": item["title"], "abstract": item["abstract"]}
        for item in records
    ]
    body = {
        "model": model,
        "prompt": f"{PROMPT}\n\nINPUT:\n{json.dumps(compact, ensure_ascii=False)}",
        "stream": False,
        "format": "json",
        "options": {"temperature": 0, "num_predict": 4000},
    }
    request = urllib.request.Request(
        f"{url.rstrip('/')}/api/generate",
        data=json.dumps(body).encode("utf-8"),
        headers={"Content-Type": "application/json"},
        method="POST",
    )
    for attempt in range(retries + 1):
        try:
            with urllib.request.urlopen(request, timeout=timeout) as response:
                result = json.loads(response.read().decode("utf-8"))
            decoded = json.loads(result["response"])
            if isinstance(decoded, list):
                parsed = decoded
            elif isinstance(decoded, dict) and isinstance(decoded.get("papers"), list):
                parsed = decoded["papers"]
            elif isinstance(decoded, dict) and isinstance(decoded.get("results"), list):
                parsed = decoded["results"]
            elif isinstance(decoded, dict) and decoded.get("arxiv_id"):
                parsed = [decoded]
            else:
                raise ValueError("response lacks a recognized result array")
            expected = [item["arxiv_id"] for item in records]
            by_id = {str(item.get("arxiv_id")): item for item in parsed}
            if any(paper_id not in by_id for paper_id in expected):
                raise ValueError("response ids do not match input")
            return [{**by_id[paper_id], "source": "model"} for paper_id in expected]
        except (
            urllib.error.HTTPError,
            urllib.error.URLError,
            TimeoutError,
            json.JSONDecodeError,
            KeyError,
            ValueError,
        ):
            if attempt >= retries:
                raise
            time.sleep(2 ** attempt)
    raise RuntimeError("unreachable retry state")


def heuristic_screen(record: dict) -> dict:
    title = record["title"].lower()
    text = f"{record['title']} {record['abstract']}".lower()
    title_memory = bool(re.search(r"memory|memori|remember|forget|stateful", title))
    agent_memory = bool(
        re.search(r"agent(?:ic)? memory|memory (?:for|in|of) .*agent|long-term .*agent", text)
    )
    hardware_only = bool(re.search(r"gpu memory|memory footprint|memory bandwidth|memory efficient", title))
    if title_memory and not hardware_only:
        relevance = "core"
    elif agent_memory:
        relevance = "supporting"
    else:
        relevance = "peripheral"

    if re.search(r"attack|poison|privacy|security|trust|provenance|leak", text):
        problem = "security_privacy"
    elif re.search(r"benchmark|evaluation|diagnos|measure", title):
        problem = "evaluation"
    elif re.search(r"update|forget|consolidat|stale|supersed|write", text):
        problem = "update_forgetting"
    elif re.search(r"skill|experience|reflect|self-evolv|continual learn", text):
        problem = "experience_learning"
    elif re.search(r"execution state|task state|workflow state|long-horizon", text):
        problem = "execution_state"
    elif re.search(r"graph|hierarch|represent|structure|database|document", text):
        problem = "representation"
    elif re.search(r"retriev|search|recall|access", text):
        problem = "access"
    elif re.search(r"multi-agent|shared memory|collective memory", text):
        problem = "shared_memory"
    elif re.search(r"cost|latency|token|efficient|compress", text):
        problem = "systems_cost"
    else:
        problem = "application" if relevance != "peripheral" else "unrelated"
    return {
        "arxiv_id": record["arxiv_id"],
        "relevance": relevance,
        "problem": problem,
        "reason": "deterministic fallback after model-format failure",
        "source": "heuristic",
    }


def normalize_analysis(record: dict, analysis: dict | None) -> dict:
    """Constrain model output to the documented screening schema."""
    fallback = heuristic_screen(record)
    analysis = analysis if isinstance(analysis, dict) else {}
    relevance = str(analysis.get("relevance", "")).strip().lower()
    if relevance not in RELEVANCE_VALUES:
        relevance = fallback["relevance"]

    problem = str(analysis.get("problem", "")).strip().lower()
    problem = PROBLEM_ALIASES.get(problem, problem)
    if problem not in PROBLEM_VALUES:
        problem = fallback["problem"]

    reason = str(analysis.get("reason", "")).strip() or fallback["reason"]
    source = str(analysis.get("source", "")).strip().lower()
    if source not in {"model", "heuristic"}:
        source = "model" if analysis else "heuristic"
    return {
        "arxiv_id": record["arxiv_id"],
        "relevance": relevance,
        "problem": problem,
        "reason": reason,
        "source": source,
    }


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--input",
        type=Path,
        default=ROOT / "data" / "research" / "agent-memory-expanded.json",
    )
    parser.add_argument(
        "--output",
        type=Path,
        default=ROOT / "data" / "research" / "agent-memory-screened.json",
    )
    parser.add_argument("--ollama-url", default="http://192.168.1.10:11434")
    parser.add_argument("--model", default="ChatGPT-5.6:Luna")
    parser.add_argument("--batch-size", type=int, default=24)
    parser.add_argument("--timeout", type=int, default=180)
    parser.add_argument("--retries", type=int, default=2)
    parser.add_argument("--refresh", action="store_true")
    args = parser.parse_args()
    args.input = args.input.resolve()
    args.output = args.output.resolve()

    source = json.loads(args.input.read_text(encoding="utf-8"))
    records = source["papers"]
    completed: dict[str, dict] = {}
    if args.output.exists() and not args.refresh:
        previous = json.loads(args.output.read_text(encoding="utf-8"))
        completed = {
            item["arxiv_id"]: {
                **item,
                "screen": normalize_analysis(item, item.get("screen")),
            }
            for item in previous.get("papers", [])
        }

    pending = [item for item in records if item["arxiv_id"] not in completed]
    for offset in range(0, len(pending), args.batch_size):
        batch = pending[offset : offset + args.batch_size]
        try:
            screened = generate(
                args.ollama_url, args.model, batch, timeout=args.timeout, retries=args.retries
            )
        except Exception:
            screened = []
            for record in batch:
                try:
                    screened.extend(
                        generate(
                            args.ollama_url,
                            args.model,
                            [record],
                            timeout=args.timeout,
                            retries=args.retries,
                        )
                    )
                except Exception:
                    screened.append(heuristic_screen(record))
        for record, analysis in zip(batch, screened, strict=True):
            completed[record["arxiv_id"]] = {
                **record,
                "screen": normalize_analysis(record, analysis),
            }
        ordered = [completed[item["arxiv_id"]] for item in records if item["arxiv_id"] in completed]
        write_json(
            args.output,
            {
                "generated_at": date.today().isoformat(),
                "model": args.model,
                "source": str(args.input.relative_to(ROOT)),
                "total": len(records),
                "completed": len(ordered),
                "papers": ordered,
            },
        )
        print(f"screened {len(ordered)}/{len(records)}", flush=True)

    # Rewrites an already complete file after schema normalization as well.
    ordered = [completed[item["arxiv_id"]] for item in records if item["arxiv_id"] in completed]
    write_json(
        args.output,
        {
            "generated_at": date.today().isoformat(),
            "model": args.model,
            "source": str(args.input.relative_to(ROOT)),
            "total": len(records),
            "completed": len(ordered),
            "papers": ordered,
        },
    )
    return 0


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