#!/usr/bin/env python3
"""Second-round live audit (2026-05-15 evening).

Different papers, different domains. Exercises the same five
subsystems as ``run_full_pipeline_audit.py`` plus every bug-fix from
the 2026-05-15 audit-bug-fix batch:

  #1 ProvenanceStore.init_db() standalone
  #2 SourceReference.authors: list[str] with validator
  #3 _fetch_seed_work arXiv chain (title.search via arXiv API)
  #4 PaperSource.PUBMED / ARXIV / OPENALEX / CROSSREF
  #5 BudgetTracker(max_tokens=..., max_cost_usd=...) kwargs
  #6 KBRouteHit destructuring
  #7 (new) cite-graph end-to-end on a fresh arXiv DOI via the
     title.search chain (replaces the broken ids.arxiv filter)

Three real papers from three different domains:
  - CRISPR-Cas9 (Jinek 2012, biomedical)            10.1126/science.1225829
  - LIGO gravitational waves (Abbott 2016, physics) 10.1103/PhysRevLett.116.061102
  - GPT-3 (Brown 2020, arXiv-only ML preprint)      10.48550/arXiv.2005.14165

Run:
    python tests/audit/run_second_round_audit.py
"""
from __future__ import annotations

import asyncio
import json
import os
import sys
import tempfile
import time
import traceback
from datetime import datetime
from pathlib import Path
from typing import Any

ROOT = Path(__file__).resolve().parents[2]
sys.path.insert(0, str(ROOT / "src"))


ARTICLES = [
    {
        "id": "crispr_cas9",
        "name": "A Programmable Dual-RNA-Guided DNA Endonuclease (CRISPR-Cas9)",
        "doi": "10.1126/science.1225829",
        "domain": "biomedical / chemistry",
        "expected_year": 2012,
    },
    {
        "id": "ligo_gw150914",
        "name": "Observation of Gravitational Waves from a Binary Black Hole Merger",
        "doi": "10.1103/PhysRevLett.116.061102",
        "domain": "physics",
        "expected_year": 2016,
    },
    {
        "id": "gpt3",
        "name": "Language Models are Few-Shot Learners (GPT-3, arXiv-only)",
        "doi": "10.48550/arXiv.2005.14165",
        "domain": "ML / NLP (arXiv-only)",
        "expected_year": 2020,
    },
]

_audit_repo = Path.home() / "git" / "research-tools-audit"
_default_results = _audit_repo / "results" / "perspicacite-current" if _audit_repo.is_dir() else ROOT / "tests" / "audit" / "results"
RESULTS_DIR = Path(os.environ["AUDIT_RESULTS_DIR"]).expanduser() if os.environ.get("AUDIT_RESULTS_DIR") else _default_results
RESULTS_DIR.mkdir(parents=True, exist_ok=True)


def section(label: str) -> None:
    print(f"\n{'=' * 4} {label} {'=' * max(0, 60 - len(label))}")


def report(msg: str) -> None:
    print(f"  {msg}")


# ---------------------------------------------------------------------------
# Phase 1 — Fresh-paper DOI → OpenAlex resolution
# (validates fix #3 + #7: arXiv DOIs now resolve via title.search chain)
# ---------------------------------------------------------------------------

async def audit_doi_resolution(findings: dict[str, Any]) -> None:
    section("Phase 1: DOI → OpenAlex resolution (fix #3 + #7)")
    import httpx

    from perspicacite.pipeline.snowball import _fetch_seed_work, openalex_id_for_doi

    out: dict[str, Any] = {}
    findings["doi_resolution"] = out
    headers = {
        "User-Agent": "perspicacite-audit (louisfelix.nothias@gmail.com)",
    }
    async with httpx.AsyncClient() as client:
        for art in ARTICLES:
            art_out: dict[str, Any] = {}
            out[art["id"]] = art_out
            t0 = time.perf_counter()
            oa_id = await openalex_id_for_doi(client, art["doi"], headers=headers)
            seed_work = await _fetch_seed_work(client, art["doi"], headers)
            elapsed = round(time.perf_counter() - t0, 2)
            art_out["openalex_id"] = oa_id
            art_out["seconds"] = elapsed
            if seed_work:
                art_out["resolved_title"] = (
                    seed_work.get("title") or seed_work.get("display_name")
                )
                art_out["cited_by_count"] = seed_work.get("cited_by_count")
                art_out["publication_year"] = seed_work.get("publication_year")
                art_out["status"] = "ok"
            else:
                art_out["status"] = "miss"
            report(
                f"[{art['id']:>20}] {art['domain']:<24} doi → "
                f"{oa_id or '<miss>'}  ({elapsed}s)"
            )
            if seed_work:
                report(
                    f"  title='{(art_out['resolved_title'] or '')[:60]}' "
                    f"cited_by={art_out['cited_by_count']}"
                )


# ---------------------------------------------------------------------------
# Phase 2 — Bug-fix validation suite
# ---------------------------------------------------------------------------

async def audit_bug_fixes(findings: dict[str, Any]) -> None:
    section("Phase 2: Bug-fix validation (audit batch fixes #1-#7)")
    bug_fixes: dict[str, Any] = {}
    findings["bug_fixes"] = bug_fixes

    # ---- Fix #1: ProvenanceStore.init_db() standalone ---------------------
    from perspicacite.provenance.store import ProvenanceStore

    tmp = Path(tempfile.mkdtemp(prefix="audit2_prov_"))
    db = tmp / "standalone.sqlite"
    sidecar = tmp / "sidecar"
    store = ProvenanceStore(db_path=db, sidecar_dir=sidecar)
    try:
        await store.init_db()
        await store.save({
            "message_id": "round2-msg-1",
            "conversation_id": "round2-conv-1",
            "rag_mode": "basic",
            "request_params": {"q": "round 2"},
            "llm_calls": [{"provider": "fake", "model": "fake-1",
                           "prompt_tokens": 5, "completion_tokens": 5}],
        })
        rec = await store.get_for_message("round2-msg-1")
        bug_fixes["fix_1_provenance_init_db"] = {
            "status": "PASS" if rec is not None else "FAIL",
            "round_trip_ok": rec is not None,
        }
        report(f"  fix#1 ProvenanceStore standalone: {'PASS' if rec else 'FAIL'}")
    except Exception as exc:
        bug_fixes["fix_1_provenance_init_db"] = {
            "status": "FAIL", "error": f"{type(exc).__name__}: {exc}"
        }
        report(f"  fix#1 ProvenanceStore standalone: FAIL — {exc}")

    # ---- Fix #2: SourceReference.authors: list[str] -----------------------
    from perspicacite.models.rag import SourceReference

    try:
        # List input
        s1 = SourceReference(title="T", authors=["Jumper", "Evans", "Pritzel"])
        # Legacy str input
        s2 = SourceReference(title="T", authors="Alice, Bob")
        # None
        s3 = SourceReference(title="T", authors=None)
        # to_citation: year=None → no trailing ", year" (empty string)
        s4 = SourceReference(title="T", authors=["Jumper"], year=2021)
        ok = (
            s1.authors == ["Jumper", "Evans", "Pritzel"]
            and s2.authors == ["Alice", "Bob"]
            and s3.authors == []
            and s1.to_citation() == "[Jumper et al.]"   # year=None → no comma
            and s4.to_citation() == "[Jumper, 2021]"
        )
        bug_fixes["fix_2_source_reference_authors_list"] = {
            "status": "PASS" if ok else "FAIL",
            "list_input": s1.authors,
            "str_input": s2.authors,
            "none_input": s3.authors,
        }
        report(f"  fix#2 SourceReference.authors list: {'PASS' if ok else 'FAIL'}")
    except Exception as exc:
        bug_fixes["fix_2_source_reference_authors_list"] = {
            "status": "FAIL", "error": f"{type(exc).__name__}: {exc}"
        }
        report(f"  fix#2 SourceReference.authors list: FAIL — {exc}")

    # ---- Fix #4: PaperSource new enum values ------------------------------
    from perspicacite.models.papers import PaperSource

    try:
        new_values = [
            PaperSource.OPENALEX.value, PaperSource.PUBMED.value,
            PaperSource.ARXIV.value, PaperSource.CROSSREF.value,
        ]
        # Round-trip from string (chroma metadata path)
        round_trips = [PaperSource(v) for v in new_values]
        ok = (
            new_values == ["openalex", "pubmed", "arxiv", "crossref"]
            and all(rt.value == v for rt, v in zip(round_trips, new_values))
        )
        bug_fixes["fix_4_paper_source_enum"] = {
            "status": "PASS" if ok else "FAIL",
            "new_values": new_values,
        }
        report(f"  fix#4 PaperSource enum extension: {'PASS' if ok else 'FAIL'}")
    except Exception as exc:
        bug_fixes["fix_4_paper_source_enum"] = {
            "status": "FAIL", "error": f"{type(exc).__name__}: {exc}"
        }
        report(f"  fix#4 PaperSource enum extension: FAIL — {exc}")

    # ---- Fix #5: BudgetTracker max_tokens / max_cost_usd ------------------
    from perspicacite.llm.budget import BudgetExceededError, BudgetTracker

    try:
        t = BudgetTracker(max_tokens=1000, max_cost_usd=1.0)
        assert t.max_usd == 1.0, f"max_usd alias failed: {t.max_usd}"
        # Record below cap → no raise
        t.record(provider="claude_cli", model="*", input_tokens=100, output_tokens=100)
        # Record above cap → raises
        raised = False
        try:
            t.record(provider="claude_cli", model="*",
                     input_tokens=600, output_tokens=300)
        except BudgetExceededError:
            raised = True
        bug_fixes["fix_5_budget_tracker_kwargs"] = {
            "status": "PASS" if raised else "FAIL",
            "tokens_total": t.tokens_in + t.tokens_out,
            "raises_at_cap": raised,
        }
        report(f"  fix#5 BudgetTracker kwargs: {'PASS' if raised else 'FAIL'}")
    except Exception as exc:
        bug_fixes["fix_5_budget_tracker_kwargs"] = {
            "status": "FAIL", "error": f"{type(exc).__name__}: {exc}"
        }
        report(f"  fix#5 BudgetTracker kwargs: FAIL — {exc}")

    # ---- Fix #6: KBRouteHit destructuring ---------------------------------
    from perspicacite.rag.kb_router import (
        KBRouteHit,
        _bm25_cache_clear,
        route_kbs,
    )

    try:
        _bm25_cache_clear()
        hit = KBRouteHit(kb_name="crispr", score=0.9)
        name, score = hit
        # Now try a real route_kbs path
        contexts = {
            "crispr_kb": "CRISPR Cas9 genome editing endonuclease",
            "physics_kb": "gravitational waves general relativity black hole",
            "ml_kb":      "transformer language model neural network",
        }
        ranked = list(route_kbs(
            query="dual RNA endonuclease genome editing",
            kb_contexts=contexts, top_k=3,
        ))
        first_name, first_score = ranked[0]
        ok = (
            name == "crispr" and score == 0.9
            and first_name == "crispr_kb"
            and isinstance(first_score, float)
        )
        bug_fixes["fix_6_kb_route_hit_iter"] = {
            "status": "PASS" if ok else "FAIL",
            "ranked": [(n, round(s, 3)) for n, s in ranked],
        }
        report(
            f"  fix#6 KBRouteHit destructure: {'PASS' if ok else 'FAIL'} "
            f"top={first_name}@{round(first_score, 3)}"
        )
    except Exception as exc:
        bug_fixes["fix_6_kb_route_hit_iter"] = {
            "status": "FAIL", "error": f"{type(exc).__name__}: {exc}"
        }
        report(f"  fix#6 KBRouteHit destructure: FAIL — {exc}")


# ---------------------------------------------------------------------------
# Phase 3 — Full ingest cycle for the three fresh papers
# (OpenAlex fetch → Paper construction → chunking → capsule artifacts)
# ---------------------------------------------------------------------------

async def audit_ingest_cycle(findings: dict[str, Any]) -> None:
    section("Phase 3: Fresh-paper ingest cycle (Lab + ASB)")
    import httpx

    from perspicacite.models.papers import Paper, PaperSource
    from perspicacite.pipeline.capsule_builder import (
        capsule_dir_for,
        write_blocks,
        write_metadata,
        write_resources,
    )
    from perspicacite.pipeline.chunking_dispatch import chunk_document

    out: dict[str, Any] = {}
    findings["ingest_cycle"] = out
    capsule_root = Path(tempfile.mkdtemp(prefix="audit2_capsule_"))
    out["capsule_root"] = str(capsule_root)

    from types import SimpleNamespace
    chunk_config = SimpleNamespace(
        knowledge_base=SimpleNamespace(
            chunk_size=512, chunk_overlap=64,
            embedding_model="text-embedding-3-small",
            use_two_pass=True, default_top_k=10,
        ),
    )

    headers = {
        "User-Agent": "perspicacite-audit (louisfelix.nothias@gmail.com)",
    }
    async with httpx.AsyncClient() as client:
        from perspicacite.pipeline.snowball import _fetch_seed_work
        for art in ARTICLES:
            art_out: dict[str, Any] = {}
            out[art["id"]] = art_out
            seed = await _fetch_seed_work(client, art["doi"], headers)
            if seed is None:
                art_out["status"] = "no_seed"
                report(f"[{art['id']}] could not resolve seed work — skip")
                continue
            inverted = seed.get("abstract_inverted_index") or {}
            positions: dict[int, str] = {}
            for word, idxs in inverted.items():
                for i in idxs:
                    positions[i] = word
            abstract = " ".join(positions[i] for i in sorted(positions))
            art_out["abstract_len"] = len(abstract)
            art_out["cited_by_count"] = seed.get("cited_by_count")

            # Pick a domain-correct PaperSource value (uses the new enum
            # values from fix #4).
            if art["id"] == "gpt3":
                source = PaperSource.ARXIV
            elif art["id"] == "crispr_cas9":
                source = PaperSource.CROSSREF
            else:
                source = PaperSource.OPENALEX

            paper = Paper(
                id=f"doi:{art['doi']}",
                title=seed.get("title") or art["name"],
                abstract=abstract,
                source=source,
                doi=art["doi"],
                year=seed.get("publication_year"),
            )
            art_out["paper_source"] = paper.source.value

            # Chunk dispatch
            if abstract:
                chunks = await chunk_document(
                    abstract, paper, content_type="text",
                    language=None, config=chunk_config,
                )
                art_out["chunk_count"] = len(chunks)
            else:
                art_out["chunk_count"] = 0

            # Capsule artifacts
            cdir = capsule_dir_for(paper, root=capsule_root)
            cdir.mkdir(parents=True, exist_ok=True)
            write_metadata(cdir, paper=paper, producer_version="audit2-0.1")
            blocks_n = write_blocks(cdir, text=abstract or "")
            res_n = write_resources(cdir, text=abstract or "")
            art_out["blocks_written"] = blocks_n
            art_out["resources_written"] = res_n
            art_out["status"] = "ok"
            report(
                f"[{art['id']:>20}] source={paper.source.value:<9} "
                f"chunks={art_out['chunk_count']} blocks={blocks_n} "
                f"resources={res_n}"
            )


# ---------------------------------------------------------------------------
# Phase 4 — Cite-graph end-to-end on the three fresh DOIs (no --openalex-id)
# Validates that arXiv DOIs now work via the title.search chain in production.
# ---------------------------------------------------------------------------

async def audit_cite_graph(findings: dict[str, Any]) -> None:
    section("Phase 4: Cite-graph by DOI (live, no --openalex-id)")
    import httpx

    from perspicacite.pipeline.snowball import _fetch_seed_work, fetch_cited_by_works

    out: dict[str, Any] = {}
    findings["cite_graph"] = out
    headers = {
        "User-Agent": "perspicacite-audit (louisfelix.nothias@gmail.com)",
    }
    async with httpx.AsyncClient() as client:
        for art in ARTICLES:
            art_out: dict[str, Any] = {}
            out[art["id"]] = art_out
            t0 = time.perf_counter()
            seed = await _fetch_seed_work(client, art["doi"], headers)
            if seed is None:
                art_out["status"] = "no_seed"
                report(f"[{art['id']}] SEED MISS — 0 hits")
                continue
            cites = await fetch_cited_by_works(
                client, seed_work=seed, max_results=10, headers=headers,
            )
            elapsed = round(time.perf_counter() - t0, 2)
            art_out["status"] = "ok"
            art_out["seconds"] = elapsed
            art_out["hit_count"] = len(cites)
            art_out["sample_titles"] = [
                (w.get("title") or w.get("display_name") or "")[:80]
                for w in cites[:3]
            ]
            report(
                f"[{art['id']:>20}] {len(cites):>2} cite-graph hits "
                f"in {elapsed}s"
            )


# ---------------------------------------------------------------------------
# Phase 4b — Multi-chunk path on real long-form text
# (Round-1 audit reported "chunks=1" everywhere because OpenAlex abstracts
# are < chunk_size. This leg exercises the multi-chunk path on a real
# 35-KB README so we don't ship blind on it.)
# ---------------------------------------------------------------------------

_MULTI_CHUNK_FIXTURES = [
    {
        "id": "alphafold_readme",
        "url": "https://raw.githubusercontent.com/google-deepmind/alphafold/main/README.md",
        "doi": "10.1038/s41586-021-03819-2",  # for capsule path naming only
    },
    {
        "id": "transformers_root_readme",
        "url": "https://raw.githubusercontent.com/huggingface/transformers/main/README.md",
        "doi": "10.48550/arXiv.2005.11401",
    },
]


async def audit_multi_chunk(findings: dict[str, Any]) -> None:
    section("Phase 4b: Multi-chunk path on real long-form text")
    from types import SimpleNamespace

    import httpx

    from perspicacite.models.papers import Paper, PaperSource
    from perspicacite.pipeline.chunking_dispatch import chunk_document

    out: dict[str, Any] = {}
    findings["multi_chunk"] = out

    chunk_config = SimpleNamespace(
        knowledge_base=SimpleNamespace(
            chunk_size=512, chunk_overlap=64,
            embedding_model="text-embedding-3-small",
            use_two_pass=True, default_top_k=10,
        ),
    )

    async with httpx.AsyncClient(follow_redirects=True) as client:
        for fix in _MULTI_CHUNK_FIXTURES:
            fix_out: dict[str, Any] = {}
            out[fix["id"]] = fix_out
            try:
                resp = await client.get(fix["url"], timeout=20.0)
            except httpx.HTTPError as exc:
                fix_out["status"] = f"fetch_error:{exc}"
                report(f"[{fix['id']}] fetch error: {exc}")
                continue
            if resp.status_code != 200:
                fix_out["status"] = f"http_{resp.status_code}"
                report(f"[{fix['id']}] HTTP {resp.status_code}")
                continue
            body = resp.text
            fix_out["body_len"] = len(body)
            paper = Paper(
                id=f"doi:{fix['doi']}", title=fix["id"], abstract="",
                source=PaperSource.OPENALEX, doi=fix["doi"],
            )
            chunks = await chunk_document(
                body, paper, content_type="text",
                language=None, config=chunk_config,
            )
            fix_out["chunk_count"] = len(chunks)
            fix_out["status"] = "ok" if len(chunks) > 1 else "ONE_CHUNK_UNEXPECTED"
            fix_out["first_chunk_preview"] = chunks[0].text[:80] if chunks else ""
            report(
                f"[{fix['id']:>24}] body={len(body):>6} → {len(chunks)} chunks "
                f"({fix_out['status']})"
            )


# ---------------------------------------------------------------------------
# Phase 5 — RAGResponse assembly + StreamEvent factories
# ---------------------------------------------------------------------------

def audit_response_assembly(findings: dict[str, Any]) -> None:
    section("Phase 5: RAGResponse + StreamEvent factories (Manuscript)")
    from perspicacite.models.rag import (
        FigureRef,
        RAGMode,
        RAGRequest,
        RAGResponse,
        SourceReference,
        StreamEvent,
    )
    out: dict[str, Any] = {}
    findings["response_assembly"] = out

    sources = [
        SourceReference(
            title="A Programmable Dual-RNA-Guided DNA Endonuclease",
            authors=["Jinek", "Chylinski", "Fonfara", "Hauer", "Doudna", "Charpentier"],
            year=2012, doi="10.1126/science.1225829", relevance_score=0.95,
        ),
        SourceReference(
            title="Observation of Gravitational Waves",
            authors=["Abbott et al."],
            year=2016, doi="10.1103/PhysRevLett.116.061102", relevance_score=0.91,
        ),
    ]
    request = RAGRequest(
        query="What is the molecular mechanism behind CRISPR-Cas9 cleavage?",
        kb_names=["crispr_kb", "biomed_kb"],
        mode=RAGMode.ADVANCED,
    )
    resp = RAGResponse(
        answer="Cas9 uses a guide RNA to direct cleavage of dsDNA via its HNH and RuvC nuclease domains.",
        sources=sources,
        mode=RAGMode.ADVANCED,
        iterations=2,
        tokens_used=850,
        figures=[FigureRef(id="f1", paper_id="doi:10.1126/science.1225829", label="Figure 1")],
        code_excerpts=[],
    )
    out["sources_len"] = len(resp.sources)
    out["sources_with_list_authors"] = sum(
        1 for s in resp.sources if isinstance(s.authors, list) and len(s.authors) > 1
    )
    out["citations"] = [s.to_citation() for s in resp.sources]
    out["kb_names_passthrough"] = request.kb_names

    # StreamEvent factories
    out["events_built"] = []
    for ev in (
        StreamEvent.status("running"),
        StreamEvent.source(sources[0]),
        StreamEvent.figure_ref({"id": "f1", "paper_id": "p", "thumbnail_b64": "iVBOR..."}),
        StreamEvent.done(conversation_id="c1", tokens_used=850,
                         mode="advanced", iterations=2),
    ):
        out["events_built"].append(ev.event)
    report(f"  sources={out['sources_len']} list-authors={out['sources_with_list_authors']}")
    report(f"  citations[0] = {out['citations'][0]}")
    report(f"  StreamEvents emitted: {out['events_built']}")


# ---------------------------------------------------------------------------
# Driver
# ---------------------------------------------------------------------------

async def main() -> int:
    started = datetime.utcnow()
    findings: dict[str, Any] = {
        "started_at": started.isoformat() + "Z",
        "git_sha": _git_sha(),
        "articles": ARTICLES,
    }

    try:
        await audit_doi_resolution(findings)
    except Exception as exc:
        findings["doi_resolution_error"] = traceback.format_exc()
        report(f"PHASE 1 ERROR: {exc}")
    try:
        await audit_bug_fixes(findings)
    except Exception as exc:
        findings["bug_fixes_error"] = traceback.format_exc()
        report(f"PHASE 2 ERROR: {exc}")
    try:
        await audit_ingest_cycle(findings)
    except Exception as exc:
        findings["ingest_cycle_error"] = traceback.format_exc()
        report(f"PHASE 3 ERROR: {exc}")
    try:
        await audit_cite_graph(findings)
    except Exception as exc:
        findings["cite_graph_error"] = traceback.format_exc()
        report(f"PHASE 4 ERROR: {exc}")
    try:
        await audit_multi_chunk(findings)
    except Exception as exc:
        findings["multi_chunk_error"] = traceback.format_exc()
        report(f"PHASE 4b ERROR: {exc}")
    try:
        audit_response_assembly(findings)
    except Exception as exc:
        findings["response_assembly_error"] = traceback.format_exc()
        report(f"PHASE 5 ERROR: {exc}")

    findings["finished_at"] = datetime.utcnow().isoformat() + "Z"
    ts = started.strftime("%Y%m%d-%H%M%S")
    json_path = RESULTS_DIR / f"second-round-audit-{ts}.json"
    md_path = RESULTS_DIR / f"second-round-audit-{ts}.md"
    json_path.write_text(json.dumps(findings, indent=2, default=str))
    md_path.write_text(_render_md(findings))
    print(f"\nJSON: {json_path.relative_to(ROOT)}")
    print(f"MD:   {md_path.relative_to(ROOT)}")
    return 0


def _git_sha() -> str:
    import subprocess
    try:
        return subprocess.check_output(
            ["git", "rev-parse", "--short", "HEAD"], cwd=ROOT,
        ).decode().strip()
    except Exception:
        return "<unknown>"


def _render_md(findings: dict[str, Any]) -> str:
    lines = [
        f"# Second-round audit — {findings['started_at']}",
        f"\nGit SHA: `{findings['git_sha']}`\n",
        "## Articles\n",
    ]
    for art in findings["articles"]:
        lines.append(f"- **{art['id']}** ({art['domain']}): `{art['doi']}`")
    for phase, blob in findings.items():
        if phase in {"started_at", "finished_at", "git_sha", "articles"}:
            continue
        lines.append(f"\n## {phase}\n")
        lines.append("```json")
        lines.append(json.dumps(blob, indent=2, default=str)[:8000])
        lines.append("```")
    return "\n".join(lines)


if __name__ == "__main__":
    sys.exit(asyncio.run(main()))
