"""
extractor.py - Stage 2: per-sheet construction object extraction.

One vision call per page produces a "sheet object manifest": the discipline
plus a list of typed, grounded construction objects (each with verbatim
source_text and a location_key). Output feeds Stage 3 normalization and
Stage 4 clustering.

Grounding guard (anti-hallucination): an object is dropped if its primary
value contains numeric claims not present in its source_text. Graphical
objects (graphical_basis set, no source_text) are allowed through.
"""

import re
import threading
from concurrent.futures import ThreadPoolExecutor, as_completed
from typing import List, Dict, Optional

from backend import config
from backend.llm import call_json
from backend.prompts import (
    EXTRACTOR_SYSTEM_PROMPT,
    EXTRACTOR_USER_INSTRUCTION,
    TEXT_STRUCTURING_SYSTEM_PROMPT,
    TEXT_STRUCTURING_USER_INSTRUCTION,
    DISCIPLINE_PREFIXES,
    ATTRIBUTE_VOCAB,
)
from backend.text_coverage import (
    _norm,
    fallback_objects,
    merge_objects,
    recover_sheet_number,
    text_coverage,
)

# prefix (upper) -> discipline, longest-prefix-first for greedy matching
_PREFIX_TO_DISCIPLINE = sorted(
    ((p.upper(), disc) for disc, ps in DISCIPLINE_PREFIXES.items() for p in ps),
    key=lambda kv: -len(kv[0]),
)

_LETTERS_RE = re.compile(r"^[A-Za-z]+")
_DIGITS_RE = re.compile(r"\d+")
# "Room 124", "RM 124A", "SPACE 12" -> "124" / "124A" / "12". Conservative: an
# explicit room/rm/space keyword followed by a 2-4 digit number (+ optional letter).
_ROOM_RE = re.compile(r"(?i)\b(?:room|rm|space)\s*#?\s*(\d{2,4}[A-Za-z]?)\b")


def _room_from_text(text: Optional[str]) -> Optional[str]:
    """Best-effort room id from free text when location_key.room is missing."""
    if not text:
        return None
    m = _ROOM_RE.search(str(text))
    return m.group(1) if m else None


def discipline_from_sheet_number(sheet_number: Optional[str]) -> Optional[str]:
    """Deterministic discipline from the sheet-number letter prefix, or None."""
    if not sheet_number:
        return None
    m = _LETTERS_RE.match(sheet_number.strip())
    if not m:
        return None
    prefix = m.group(0).upper()
    for cand, disc in _PREFIX_TO_DISCIPLINE:  # longest prefix wins
        if prefix == cand:
            return disc
    # Fall back to the single leading letter (e.g. "A2" -> "A")
    for cand, disc in _PREFIX_TO_DISCIPLINE:
        if prefix.startswith(cand) and len(cand) == 1:
            return disc
    return None


def _is_grounded(value: str, source_text: str, graphical_basis: str = "",
                 page_text: Optional[str] = None) -> bool:
    """
    Keep an object only if its primary value is supported by its source_text,
    OR it is a graphical object (has graphical_basis with no text to quote).

    - If graphical_basis is set and source_text is absent, the object is valid.
    - If the value contains digits, every distinct digit-run must appear in
      source_text (catches invented dimensions/counts/elevations).
    - Rescue tier: when page_text (the deterministic text layer) is given,
      digit-runs absent from source_text but present in the page text are
      still grounded - vision quoted imperfectly but the value is real.
    - If the value has no digits, require some alphabetic-token overlap.
    """
    # Graphical objects (no readable text on sheet) are always allowed through.
    if graphical_basis and not source_text:
        return True
    if not value or not source_text:
        return False
    value = str(value)
    src_low = source_text.lower()

    val_digits = set(_DIGITS_RE.findall(value))
    if val_digits:
        src_digits = set(_DIGITS_RE.findall(source_text))
        if val_digits.issubset(src_digits):
            return True
        if page_text:
            return val_digits.issubset(set(_DIGITS_RE.findall(page_text)))
        return False

    # No digits: text-based grounding.
    val_norm = re.sub(r"[^a-z0-9]+", " ", value.lower()).strip()
    if val_norm and val_norm in src_low:
        return True
    val_tokens = {t for t in val_norm.split() if len(t) > 2}
    src_tokens = set(re.sub(r"[^a-z0-9]+", " ", src_low).split())
    return bool(val_tokens & src_tokens)


def _primary_value(obj: Dict) -> str:
    """Extract a single string 'value' from an object for grounding checks.

    Handles both the new 'attributes' dict schema and the legacy 'value' field.
    """
    attrs = obj.get("attributes")
    if isinstance(attrs, dict) and attrs:
        return str(next(iter(attrs.values())))
    # Legacy assertion schema has a top-level 'value' field
    return (obj.get("value") or obj.get("description")
            or obj.get("name") or obj.get("tag") or "")


def _grounding_stamp(value: str, source_text: str,
                     page_text: Optional[str]) -> Optional[str]:
    """\"text_layer\" when the object survived only via the text-layer rescue
    tier (digits absent from source_text but present in the page text)."""
    if not page_text:
        return None
    val_digits = set(_DIGITS_RE.findall(str(value)))
    if not val_digits:
        return None
    if val_digits.issubset(set(_DIGITS_RE.findall(source_text))):
        return None
    if val_digits.issubset(set(_DIGITS_RE.findall(page_text))):
        return "text_layer"
    return None


def _normalize_sheet(parsed: Dict, page_number: int,
                     page_text: Optional[str] = None) -> Dict:
    """
    Validate + clean one parsed sheet result, attaching page_number and ids.

    The LLM now returns a 'sheet' header + 'objects' array (Stage 2 schema).
    We map each object to an enriched assertion dict that is backward-compatible
    with all downstream stages (clusterer, conflict_checker, etc.) while also
    carrying the new typed-object fields.
    """
    # Support both new schema {sheet:{...}, objects:[...]} and old {sheet_number, assertions:[...]}
    sheet_meta = parsed.get("sheet") or {}
    sheet_number = sheet_meta.get("sheet_number") or parsed.get("sheet_number")
    sheet_title  = sheet_meta.get("sheet_title")  or parsed.get("sheet_title")
    level        = sheet_meta.get("level")         or parsed.get("level")
    scale        = sheet_meta.get("scale")         or parsed.get("scale")
    drawing_type = sheet_meta.get("drawing_type")

    raw_discipline = sheet_meta.get("discipline") or parsed.get("discipline")
    discipline = raw_discipline or discipline_from_sheet_number(sheet_number)
    if not discipline or discipline == "General":
        deduced = discipline_from_sheet_number(sheet_number)
        if deduced:
            discipline = deduced
    discipline = discipline or "Unknown"

    # Accept both new 'objects' and legacy 'assertions' keys
    raw_objects = parsed.get("objects") or parsed.get("assertions") or []
    clean: List[Dict] = []
    dropped = 0
    rescued = 0
    unverified = 0
    page_norm = _norm(page_text) if page_text else ""

    for idx, obj in enumerate(raw_objects):
        if not isinstance(obj, dict):
            continue

        source_text    = obj.get("source_text") or ""
        graphical_basis = obj.get("graphical_basis") or ""

        # Derive a primary value for the grounding check
        primary_val = _primary_value(obj)

        if not _is_grounded(primary_val, source_text, graphical_basis,
                            page_text=page_text):
            dropped += 1
            continue
        # Pre-set stamps (fallback/merge rungs) win; otherwise compute the
        # text-layer rescue stamp.
        grounding = obj.get("grounding") or _grounding_stamp(
            primary_val, source_text, page_text)
        if grounding == "text_layer":
            rescued += 1
        if not grounding and page_text and source_text:
            # Vision-unverified: survived the digit guard, but the quoted
            # source_text is not present in the deterministic text layer.
            # Kept and stamped - the wave-5b verifier prioritizes these.
            if _norm(str(source_text)) not in page_norm:
                grounding = "vision_unverified"
                unverified += 1

        # --- location_key: new schema is richer; map to legacy shape + extras ---
        lk = obj.get("location_key")
        if not isinstance(lk, dict):
            lk = {}

        # Room recovery: try new keys first, then fall back to text scan
        room = (lk.get("room") or lk.get("room_number")
                or _room_from_text(obj.get("name")) or _room_from_text(source_text))
        room_name = lk.get("room_name")

        assertion_id = (obj.get("object_id") or obj.get("assertion_id")
                        or f"{sheet_number or 'p' + str(page_number)}#{idx}")

        # Map object's attributes dict to legacy subject/attribute/value
        attrs = obj.get("attributes") or {}
        if isinstance(attrs, dict) and attrs:
            first_attr_key = next(iter(attrs))
            first_attr_val = str(attrs[first_attr_key])
        else:
            first_attr_key = obj.get("object_type") or "unspecified"
            first_attr_val = primary_val

        clean.append({
            # --- legacy fields (backward-compat with clusterer, conflict_checker, etc.) ---
            "id":        assertion_id,
            "subject":   (obj.get("name") or obj.get("tag") or obj.get("object_type")
                          or obj.get("subject") or ""),
            "attribute": first_attr_key,
            "value":     first_attr_val,
            "location_key": {
                "grid":      lk.get("grid"),
                "room":      room,
                "room_name": room_name,
                "level":     lk.get("level") or level,
                "tag":       lk.get("tag") or obj.get("tag"),
                "detail_reference": lk.get("detail_reference"),
                "plan_zone": lk.get("plan_zone"),
                "elevation_reference": lk.get("elevation_reference"),
            },
            "source_text": source_text,
            "confidence":  obj.get("confidence") or "medium",
            # --- new typed-object fields ---
            "object_type":       obj.get("object_type"),
            "category":          obj.get("category"),
            "object_tag":        obj.get("tag"),
            "object_name":       obj.get("name"),
            "object_description": obj.get("description"),
            "object_attributes": attrs,
            "graphical_basis":   graphical_basis or None,
            "review_uses":       obj.get("review_uses") or [],
            **({"grounding": grounding} if grounding else {}),
        })

    if dropped or rescued or unverified:
        print(f"[Extract] Page {page_number} ({sheet_number}): "
              f"dropped {dropped} ungrounded object(s)"
              + (f", rescued {rescued} via text layer" if rescued else "")
              + (f", {unverified} vision-unverified" if unverified else ""))

    unresolved = parsed.get("unresolved_items") or []

    return {
        "page_number":     page_number,
        "sheet_number":    sheet_number,
        "discipline":      discipline,
        "sheet_title":     sheet_title,
        "drawing_type":    drawing_type,
        "level":           level,
        "scale":           scale,
        "assertions":      clean,         # downstream stages read this key
        "unresolved_items": unresolved,
    }


def _text_layer_block(page: Dict) -> str:
    """
    The TEXT LAYER block appended to the extractor instruction at call sites
    (NOT a template placeholder - render() silently leaves missing keys as
    literals). Empty string when the page has no usable text layer.
    """
    text = (page.get("text_layer") or "").strip()
    if not text:
        return ""
    return ("\n\nTEXT LAYER (authoritative for alphanumeric content — trust it "
            "over the image for numbers, tags, and note text):\n"
            + text[:config.TEXT_LAYER_MAX_CHARS])


def _text_structuring_extract(page: Dict, sheet_hint: str = ""):
    """Rung 2 of the extraction ladder: text-only structuring call (no
    image). The text layer is authoritative for alphanumeric content - the
    model segments it instead of transcribing pixels, so vision misreads
    are impossible on this rung."""
    instruction = (TEXT_STRUCTURING_USER_INSTRUCTION
                   .replace("{sheet_hint}", sheet_hint or "")
                   .replace("{text_layer}",
                            (page.get("text_layer") or "")
                            [:config.TEXT_LAYER_MAX_CHARS]))
    return call_json(
        system_prompt=TEXT_STRUCTURING_SYSTEM_PROMPT,
        user_text=instruction,
        max_tokens=config.EXTRACT_MAX_TOKENS,
    )


def _extract_one(page: Dict, sheet_hint: str = "") -> Dict:
    page_text = page.get("text_layer")
    user_text = (EXTRACTOR_USER_INSTRUCTION.replace("{sheet_hint}", sheet_hint)
                 + _text_layer_block(page))
    parsed = call_json(
        system_prompt=EXTRACTOR_SYSTEM_PROMPT,
        user_text=user_text,
        images_b64=[page["base64"]],
        max_tokens=config.EXTRACT_MAX_TOKENS,
    )
    if not isinstance(parsed, dict):
        if not page_text:
            # Scanned/raster page: vision-only, keep the legacy failure shape.
            return {
                "page_number": page["page_number"],
                "sheet_number": None,
                "discipline": "Unknown",
                "sheet_title": f"Page {page['page_number']} (extraction failed)",
                "level": None,
                "scale": None,
                "assertions": [],
            }
        # Text-bearing page: climb the ladder instead of going dark.
        parsed = {"sheet": {}, "objects": []}

    sheet = _normalize_sheet(parsed, page["page_number"], page_text=page_text)
    cov = text_coverage(page_text or "", sheet["assertions"])
    sheet["coverage"] = cov

    # Rung 2: text-only structuring when coverage is below floor. MERGE,
    # never replace - vision objects (graphical_basis content exists only
    # in the image) are kept; the text pass fills what vision missed.
    if (page_text and config.EXTRACT_TEXT_RETRY_ENABLED
            and cov["ratio"] < config.EXTRACT_COVERAGE_FLOOR):
        print(f"[Extract] Page {page['page_number']}: coverage "
              f"{cov['ratio']:.0%} < floor - text-only structuring pass")
        parsed2 = _text_structuring_extract(page, sheet_hint)
        if isinstance(parsed2, dict):
            sheet2 = _normalize_sheet(parsed2, page["page_number"],
                                      page_text=page_text)
            before = len(sheet["assertions"])
            sheet["assertions"] = merge_objects(sheet["assertions"],
                                                sheet2["assertions"])
            for key in ("sheet_number", "sheet_title", "discipline",
                        "level", "scale", "drawing_type"):
                if not sheet.get(key) and sheet2.get(key):
                    sheet[key] = sheet2[key]
            cov = text_coverage(page_text, sheet["assertions"])
            sheet["coverage"] = cov
            print(f"[Extract] Page {page['page_number']}: merged "
                  f"{len(sheet['assertions']) - before} text-structured "
                  f"object(s), coverage now {cov['ratio']:.0%}")

    # Rung 3: deterministic fallback - a dark text-bearing sheet is
    # impossible. Stubs are deduped against earlier rungs.
    if (page_text and config.EXTRACT_FALLBACK_ENABLED
            and cov["ratio"] < config.EXTRACT_COVERAGE_FLOOR):
        stubs = fallback_objects(page_text, page["page_number"],
                                 config.EXTRACT_FALLBACK_MAX_OBJECTS)
        stubs = _normalize_sheet({"sheet": {}, "objects": stubs},
                                 page["page_number"],
                                 page_text=page_text)["assertions"]
        before = len(sheet["assertions"])
        sheet["assertions"] = merge_objects(sheet["assertions"], stubs)
        print(f"[Extract] Page {page['page_number']}: fallback merged "
              f"{len(sheet['assertions']) - before} text-layer stub(s)")
        sheet["coverage"] = text_coverage(page_text, sheet["assertions"])

    # Identity recovery: never leave a text-bearing page sheet-less.
    if not sheet.get("sheet_number") and page_text:
        recovered = recover_sheet_number(page_text)
        if recovered:
            sheet["sheet_number"] = recovered
            sheet["discipline"] = (discipline_from_sheet_number(recovered)
                                   or sheet.get("discipline") or "Unknown")
            print(f"[Extract] Page {page['page_number']}: sheet number "
                  f"recovered from text layer -> {recovered}")

    return sheet


def extract_assertions(pages: List[Dict], on_progress=None) -> List[Dict]:
    """
    Run Stage 2 over all pages. Returns a list of per-sheet object dicts,
    ordered by page_number. The 'assertions' key on each sheet holds the
    enriched construction objects (backward-compatible with downstream stages).
    """
    print(f"[Extract] Extracting construction objects from {len(pages)} pages...")
    results: List[Dict] = []
    completed = 0
    lock = threading.Lock()

    def work(page: Dict) -> Dict:
        nonlocal completed
        out = _extract_one(page)
        with lock:
            completed += 1
            n = len(out["assertions"])
            u = len(out.get("unresolved_items") or [])
            print(f"[Extract] Page {page['page_number']} -> {out['discipline']} "
                  f"({out.get('sheet_number')}), {n} objects"
                  + (f", {u} unresolved" if u else "")
                  + f" [{completed}/{len(pages)}]")
            if on_progress:
                on_progress(completed, len(pages))
        return out

    with ThreadPoolExecutor(max_workers=config.EXTRACT_CONCURRENCY) as ex:
        futures = {ex.submit(work, p): p for p in pages}
        for fut in as_completed(futures):
            try:
                results.append(fut.result())
            except Exception as e:
                p = futures[fut]
                print(f"[Extract] Page {p['page_number']} unexpected error: {e}")

    results.sort(key=lambda s: s["page_number"])
    total = sum(len(s["assertions"]) for s in results)
    total_unresolved = sum(len(s.get("unresolved_items") or []) for s in results)
    print(f"[Extract] Done - {total} grounded objects across {len(results)} sheets"
          + (f" ({total_unresolved} unresolved items)" if total_unresolved else ""))
    return results
