"""
Stratum A v3 assembly per DESIGN_v3.md.

Source-specific extractors + content-validation gates.

Extractors:
  - Paul Graham: simple HTML->text on stable static pages
  - Stephen Wolfram: extract <article> or .entry-content only
  - arxiv: use export API, extract <summary> (abstract)

Gates (all must pass):
  G1. Word count >= 300 (>= 150 for arxiv)
  G2. Paragraph count >= 3 (each paragraph >= 40 words)
  G3. Navigation-ratio: short-lines (<=5 tokens) <= 30% of non-empty lines

Writes:
  corpus/a04...a10.txt (for documents that pass all gates)
  corpus/manifest.json (updated)
  assembly_v3_log.json (per-document: fetched, gate results, accepted/excluded)
"""

from __future__ import annotations

import hashlib
import json
import os
import re
import subprocess
import sys
import urllib.request
import urllib.error
from pathlib import Path


REPO_ROOT = Path(__file__).resolve().parents[2]
STUDY_DIR = REPO_ROOT / "fvs_eval" / "validation_study"
CORPUS_DIR = STUDY_DIR / "corpus"


# ============================================================================
# Stratum A v3 source list (frozen per DESIGN_v3 §3)
# ============================================================================

STRATUM_A_V3 = [
    {
        "slug": "a04_pg_cities_and_ambition",
        "source_type": "paul_graham",
        "url": "https://www.paulgraham.com/cities.html",
        "title": "Paul Graham essay: Cities and Ambition",
        "notes": "Personal essay by Paul Graham. Reflective analytical register.",
    },
    {
        "slug": "a05_pg_how_you_know",
        "source_type": "paul_graham",
        "url": "https://www.paulgraham.com/know.html",
        "title": "Paul Graham essay: How You Know",
        "notes": "Personal essay by Paul Graham. Reflective analytical register.",
    },
    {
        "slug": "a06_wolfram_chatgpt",
        "source_type": "wolfram",
        "url": "https://writings.stephenwolfram.com/2023/02/what-is-chatgpt-doing-and-why-does-it-work/",
        "title": "Stephen Wolfram: What Is ChatGPT Doing and Why Does It Work?",
        "notes": "Technical-reflective essay by Stephen Wolfram on LLMs.",
    },
    {
        "slug": "a07_wolfram_fundamental_physics",
        "source_type": "wolfram",
        "url": "https://writings.stephenwolfram.com/2020/04/finally-we-may-have-a-path-to-the-fundamental-theory-of-physics-and-its-beautiful/",
        "title": "Stephen Wolfram: A Path to Fundamental Physics",
        "notes": "Scientific-reflective essay by Stephen Wolfram on physics.",
    },
    {
        "slug": "a08_arxiv_gpt3",
        "source_type": "arxiv",
        "url": "https://export.arxiv.org/api/query?id_list=2005.14165",
        "title": "arxiv 2005.14165: Language Models are Few-Shot Learners (GPT-3)",
        "notes": "Academic research abstract.",
    },
    {
        "slug": "a09_arxiv_foundation_models",
        "source_type": "arxiv",
        "url": "https://export.arxiv.org/api/query?id_list=2108.07258",
        "title": "arxiv 2108.07258: On the Opportunities and Risks of Foundation Models",
        "notes": "Academic research abstract.",
    },
    {
        "slug": "a10_arxiv_scaling_laws",
        "source_type": "arxiv",
        "url": "https://export.arxiv.org/api/query?id_list=2001.08361",
        "title": "arxiv 2001.08361: Scaling Laws for Neural Language Models",
        "notes": "Academic research abstract.",
    },
]


# ============================================================================
# Source-specific extractors
# ============================================================================


def _curl(url: str, timeout: int = 30) -> str:
    result = subprocess.run(
        ["curl", "-sSL", "--max-time", str(timeout),
         "-A", "FrameCheck-ValidationStudy/1.0 (research)", url],
        capture_output=True, text=True, timeout=timeout + 15,
    )
    if result.returncode != 0 or not result.stdout:
        raise RuntimeError(f"curl failed: rc={result.returncode}")
    return result.stdout


def _replace_html_entities(text: str) -> str:
    for old, new in [
        ("&amp;", "&"), ("&lt;", "<"), ("&gt;", ">"),
        ("&quot;", '"'), ("&#39;", "'"), ("&apos;", "'"),
        ("&nbsp;", " "), ("&#x27;", "'"),
        ("&mdash;", "-"), ("&ndash;", "-"),
        ("&rsquo;", "'"), ("&lsquo;", "'"),
        ("&ldquo;", '"'), ("&rdquo;", '"'),
        ("&hellip;", "..."), ("&#x2014;", "-"),
    ]:
        text = text.replace(old, new)
    return text


def extract_paul_graham(html: str) -> str:
    """Extract Paul Graham essay text.

    paulgraham.com uses table-based layout with single <br> tags for visual
    line breaks and double <br><br> for paragraph breaks. Correct extraction
    treats single <br> as line break (within paragraph) and double <br><br>
    as paragraph break. Previous implementation joined every stripped line
    with \\n\\n which made every visual line a paragraph; fixed here.
    """
    # Remove scripts, styles, head
    html = re.sub(r"<script[^>]*>.*?</script>", "", html, flags=re.DOTALL | re.IGNORECASE)
    html = re.sub(r"<style[^>]*>.*?</style>", "", html, flags=re.DOTALL | re.IGNORECASE)
    html = re.sub(r"<head[^>]*>.*?</head>", "", html, flags=re.DOTALL | re.IGNORECASE)
    # Find body
    m = re.search(r"<body[^>]*>(.*?)</body>", html, flags=re.DOTALL | re.IGNORECASE)
    body = m.group(1) if m else html
    # Close tags for paragraph-like containers become paragraph breaks
    body = re.sub(r"</(p|div|h[1-6]|li|blockquote)>", "\n\n", body, flags=re.IGNORECASE)
    # Single <br> is a visual line break (not a paragraph break).
    # Consecutive <br><br> will produce consecutive \n\n which is then
    # normalized to a single paragraph break.
    body = re.sub(r"<(br|hr)[^>]*/?>", "\n", body, flags=re.IGNORECASE)
    # Strip remaining tags
    text = re.sub(r"<[^>]+>", "", body)
    text = _replace_html_entities(text)
    # Normalize: 2+ consecutive newlines become paragraph boundaries.
    # Within each paragraph, single newlines are collapsed to spaces.
    text = re.sub(r"\n{2,}", "\n\n", text)
    paragraphs = text.split("\n\n")
    cleaned = []
    for p in paragraphs:
        p = re.sub(r"\n+", " ", p).strip()
        p = re.sub(r" +", " ", p)
        if p:
            cleaned.append(p)
    return "\n\n".join(cleaned)


def extract_wolfram(html: str) -> str:
    """Extract Stephen Wolfram blog post text.

    Targets the <article> or <div class="entry-content"> container.
    Raises RuntimeError if container not found (treats as fetch failure).
    """
    # Remove scripts, styles, head
    html = re.sub(r"<script[^>]*>.*?</script>", "", html, flags=re.DOTALL | re.IGNORECASE)
    html = re.sub(r"<style[^>]*>.*?</style>", "", html, flags=re.DOTALL | re.IGNORECASE)
    html = re.sub(r"<head[^>]*>.*?</head>", "", html, flags=re.DOTALL | re.IGNORECASE)
    html = re.sub(r"<nav[^>]*>.*?</nav>", "", html, flags=re.DOTALL | re.IGNORECASE)
    # Remove images and image captions
    html = re.sub(r"<img[^>]+/?>", "", html, flags=re.IGNORECASE)
    html = re.sub(r"<figcaption[^>]*>.*?</figcaption>", "", html, flags=re.DOTALL | re.IGNORECASE)

    # Try to find the article content container
    # WordPress themes: often <article>...</article> or <div class="entry-content">
    container_patterns = [
        r"<article[^>]*>(.*?)</article>",
        r'<div[^>]*class=["\'][^"\']*entry-content[^"\']*["\'][^>]*>(.*?)</div>\s*<(?:/div|footer|aside)',
        r'<div[^>]*class=["\'][^"\']*post-content[^"\']*["\'][^>]*>(.*?)</div>\s*<(?:/div|footer|aside)',
        r'<main[^>]*>(.*?)</main>',
    ]
    body = None
    for pat in container_patterns:
        m = re.search(pat, html, flags=re.DOTALL | re.IGNORECASE)
        if m:
            body = m.group(1)
            break
    if body is None:
        raise RuntimeError("Wolfram extractor: no article container found")

    # Replace block tags with newlines
    body = re.sub(r"</(p|div|br|h[1-6]|li|tr|td|blockquote)>", "\n", body, flags=re.IGNORECASE)
    body = re.sub(r"<(br|hr)[^>]*/?>", "\n", body, flags=re.IGNORECASE)
    # Strip tags
    text = re.sub(r"<[^>]+>", "", body)
    text = _replace_html_entities(text)
    lines = [ln.strip() for ln in text.split("\n")]
    lines = [ln for ln in lines if ln]
    return "\n\n".join(lines)


def extract_arxiv(xml: str) -> str:
    """Extract arxiv abstract from export API Atom feed."""
    # The abstract is in <summary>...</summary>
    m = re.search(r"<summary[^>]*>(.*?)</summary>", xml, flags=re.DOTALL)
    if not m:
        raise RuntimeError("arxiv extractor: no <summary> in response")
    abstract = m.group(1).strip()
    # Replace common XML entities (arxiv abstracts use standard XML escaping)
    abstract = _replace_html_entities(abstract)
    # arxiv abstracts often have line wraps; collapse single newlines to space,
    # preserve paragraph breaks (double newlines)
    abstract = re.sub(r"(?<!\n)\n(?!\n)", " ", abstract)
    abstract = re.sub(r" +", " ", abstract)
    abstract = abstract.strip()
    return abstract


EXTRACTORS = {
    "paul_graham": extract_paul_graham,
    "wolfram": extract_wolfram,
    "arxiv": extract_arxiv,
}


# ============================================================================
# Content-validation gates (pre-declared)
# ============================================================================


def gate_G1_word_count(text: str, source_type: str) -> tuple[bool, str]:
    """Word count gate. Arxiv has lower threshold (150) because abstracts
    are typically 150-300 words."""
    threshold = 150 if source_type == "arxiv" else 300
    wc = len(text.split())
    if wc >= threshold:
        return True, f"G1 pass: {wc} words (threshold {threshold})"
    return False, f"G1 FAIL: {wc} words below threshold {threshold}"


def gate_G2_paragraphs(text: str, source_type: str) -> tuple[bool, str]:
    """Paragraph count gate. >= 3 paragraphs of >= 40 words each.
    arxiv abstracts are often one long paragraph; relax to >= 1 paragraph
    of >= 100 words."""
    paras = [p.strip() for p in re.split(r"\n\s*\n", text) if p.strip()]
    if source_type == "arxiv":
        qualifying = [p for p in paras if len(p.split()) >= 100]
        if len(qualifying) >= 1:
            return True, f"G2 pass (arxiv): {len(qualifying)} paragraph(s) >= 100 words"
        return False, f"G2 FAIL (arxiv): no paragraph >= 100 words; had {len(paras)} para(s), largest {max((len(p.split()) for p in paras), default=0)} words"
    qualifying = [p for p in paras if len(p.split()) >= 40]
    if len(qualifying) >= 3:
        return True, f"G2 pass: {len(qualifying)} paragraphs >= 40 words"
    return False, f"G2 FAIL: only {len(qualifying)} paragraph(s) >= 40 words; total paras {len(paras)}"


def gate_G3_navigation_ratio(text: str, source_type: str) -> tuple[bool, str]:
    """Navigation-ratio gate. Short lines (<=5 tokens) should be <= 30%
    of non-empty lines. Skipped for arxiv (single paragraph doesn't have lines).
    """
    if source_type == "arxiv":
        return True, "G3 skipped (arxiv has no line structure)"
    lines = [ln.strip() for ln in text.split("\n") if ln.strip()]
    if not lines:
        return False, "G3 FAIL: empty document"
    short = sum(1 for ln in lines if len(ln.split()) <= 5)
    ratio = short / len(lines)
    if ratio <= 0.30:
        return True, f"G3 pass: {short}/{len(lines)} short lines ({ratio:.0%})"
    return False, f"G3 FAIL: {short}/{len(lines)} short lines ({ratio:.0%}); exceeds 30%"


def validate(text: str, source_type: str) -> tuple[bool, list]:
    """Run all three gates. Return (pass, log_entries)."""
    g1 = gate_G1_word_count(text, source_type)
    g2 = gate_G2_paragraphs(text, source_type)
    g3 = gate_G3_navigation_ratio(text, source_type)
    log = [g1[1], g2[1], g3[1]]
    all_pass = g1[0] and g2[0] and g3[0]
    return all_pass, log


# ============================================================================
# Assembly
# ============================================================================


def _sha256(s): return hashlib.sha256(s.encode("utf-8")).hexdigest()
def _wc(s): return len(s.split())
def _now(): return subprocess.check_output(["date", "-u", "+%Y-%m-%dT%H:%M:%SZ"]).decode().strip()


def main():
    # Load existing manifest
    mpath = CORPUS_DIR / "manifest.json"
    with open(mpath) as f:
        manifest = json.load(f)
    existing_ids = {d["doc_id"] for d in manifest["documents"]}

    log = []
    accepted = []

    print("=" * 70)
    print("Stratum A v3 assembly (DESIGN_v3)")
    print("=" * 70)

    for spec in STRATUM_A_V3:
        if spec["slug"] in existing_ids:
            print(f"\n  {spec['slug']}: SKIP (already in manifest)")
            continue

        entry = {
            "slug": spec["slug"],
            "url": spec["url"],
            "source_type": spec["source_type"],
            "fetch_outcome": None,
            "extractor_outcome": None,
            "gates": [],
            "accepted": False,
            "word_count": None,
        }

        print(f"\n  {spec['slug']} ({spec['source_type']})")
        print(f"    URL: {spec['url']}")

        # Fetch
        try:
            raw = _curl(spec["url"], timeout=30)
            entry["fetch_outcome"] = f"OK ({len(raw)} bytes)"
            print(f"    fetch: OK ({len(raw)} bytes)")
        except Exception as e:
            entry["fetch_outcome"] = f"FAIL: {e}"
            print(f"    fetch: FAIL: {e}")
            log.append(entry)
            continue

        # Extract
        extractor = EXTRACTORS[spec["source_type"]]
        try:
            text = extractor(raw)
            entry["extractor_outcome"] = f"OK ({len(text)} chars, {_wc(text)} words)"
            print(f"    extract: OK ({_wc(text)} words)")
        except Exception as e:
            entry["extractor_outcome"] = f"FAIL: {e}"
            print(f"    extract: FAIL: {e}")
            log.append(entry)
            continue

        # Validate
        passed, gate_log = validate(text, spec["source_type"])
        entry["gates"] = gate_log
        for g in gate_log:
            print(f"    {g}")

        if not passed:
            entry["accepted"] = False
            print(f"    -> EXCLUDED (gate failure)")
            log.append(entry)
            continue

        # Accept: write file and add to manifest
        (CORPUS_DIR / f"{spec['slug']}.txt").write_text(text, encoding="utf-8")
        entry["accepted"] = True
        entry["word_count"] = _wc(text)
        doc_entry = {
            "doc_id": spec["slug"],
            "stratum": "A",
            "title": spec["title"],
            "source": spec["url"],
            "source_type": spec["source_type"],
            "fetched_at": _now(),
            "word_count": _wc(text),
            "sha256": _sha256(text),
            "notes": spec["notes"],
            "validation_gates": gate_log,
        }
        accepted.append(doc_entry)
        print(f"    -> ACCEPTED ({_wc(text)} words)")
        log.append(entry)

    # Update manifest
    all_docs = manifest["documents"] + accepted
    strata = {"A": 0, "B": 0, "C": 0}
    for d in all_docs:
        strata[d["stratum"]] += 1

    manifest_out = {
        "version": 5,
        "assembled_at": _now(),
        "total_documents": len(all_docs),
        "target": 30,
        "halt_threshold": 27,
        "per_stratum": strata,
        "design_notes": (
            "v5 corpus after DESIGN_v3 Stratum A re-assembly. Source-specific "
            "extractors + content-validation gates applied per DESIGN_v3.md §3-§4. "
            f"Accepted {len(accepted)} new Stratum A documents. "
            f"Total n={len(all_docs)} "
            + ("(>= halt threshold, proceeding)" if len(all_docs) >= 27
               else "(< halt threshold, HALT)")
        ),
        "documents": all_docs,
    }
    with open(mpath, "w") as f:
        json.dump(manifest_out, f, indent=2)

    # Write assembly log
    log_path = STUDY_DIR / "assembly_v3_log.json"
    with open(log_path, "w") as f:
        json.dump({
            "version": 1,
            "executed_at": _now(),
            "target_new_docs": len(STRATUM_A_V3),
            "accepted": len(accepted),
            "log": log,
        }, f, indent=2)

    print("\n" + "=" * 70)
    print("Summary")
    print("=" * 70)
    print(f"  Target new Stratum A docs: {len(STRATUM_A_V3)}")
    print(f"  Accepted: {len(accepted)}")
    print(f"  Rejected: {len(STRATUM_A_V3) - len(accepted)}")
    print(f"  Total corpus n: {len(all_docs)} (A={strata['A']}, B={strata['B']}, C={strata['C']})")
    print(f"  Halt threshold: 27")
    if len(all_docs) >= 27:
        print(f"  Status: PROCEED (n >= 27)")
    else:
        print(f"  Status: HALT (n < 27); see log for gate failures")
    print(f"\n  Assembly log: {log_path}")
    print(f"  Manifest: {mpath}")


if __name__ == "__main__":
    main()