"""Scoped extraction and orientation agents."""

import json
from typing import Dict

from backend import config
from backend.agents.base import AgentResult, AgentScope, AgentUsage, failure
from backend.llm import call_json
from backend.pipeline.extractor import (
    _normalize_sheet,
    _text_layer_block,
    discipline_from_sheet_number,
)
from backend.pipeline.sheet_index import _index_input
from backend.prompts import (
    EXTRACTOR_SYSTEM_PROMPT,
    EXTRACTOR_USER_INSTRUCTION,
    JURISDICTION_SYSTEM_PROMPT,
    JURISDICTION_USER_INSTRUCTION,
    SHEET_INDEX_SYSTEM_PROMPT,
    SHEET_INDEX_USER_INSTRUCTION,
)


# Appended to the extractor instruction on the second-chance retry. Dense plan
# sheets blow the output budget on the full schema; compact mode trades
# per-object verbosity for actually finishing the page.
_COMPACT_RETRY_SUFFIX = """
IMPORTANT - COMPACT RETRY: the first pass did not complete. Keep the SAME JSON
schema, but extract at most 40 objects, prioritizing coordination-relevant
items (equipment, fixtures, devices, keynotes, dimensions, markers/callouts,
schedule rows). Keep descriptions/attributes short; skip review_uses entries
you are unsure about. Finish the JSON - a smaller complete answer beats a
larger truncated one."""


def _wrap_bare_list(parsed, page_number: int):
    """Models sometimes skip the {sheet, objects} wrapper and return a bare
    objects array (especially after truncation repair). Accept it - the sheet
    header falls back to title-block deduction downstream."""
    if isinstance(parsed, list):
        print(f"[Extract] Page {page_number}: wrapping bare objects array "
              f"({len(parsed)} items, no sheet header)")
        return {"sheet": {}, "objects": parsed}
    return parsed


class SheetExtractorAgent:
    name = "sheet_extractor"

    def __init__(self, usage: AgentUsage) -> None:
        self.usage = usage

    def _call(self, instruction: str, page: Dict):
        return call_json(
            system_prompt=EXTRACTOR_SYSTEM_PROMPT,
            user_text=instruction,
            images_b64=[page["base64"]],
            max_tokens=config.EXTRACT_MAX_TOKENS,
            model=config.AGENT_EXTRACT_MODEL,
            usage_tracker=self.usage,
            usage_stage="agent.extract",
            reasoning_effort=config.EXTRACT_REASONING_EFFORT or None,
            reasoning_max_tokens=config.EXTRACT_REASONING_MAX_TOKENS or None,
        )

    def _text_structuring_call(self, instruction_page: Dict, sheet_hint: str):
        """Rung 2: text-only structuring pass over the page's text layer
        (no image). Recovers text content the vision pass missed."""
        from backend.prompts import (TEXT_STRUCTURING_SYSTEM_PROMPT,
                                     TEXT_STRUCTURING_USER_INSTRUCTION)
        instruction = (TEXT_STRUCTURING_USER_INSTRUCTION
                       .replace("{sheet_hint}", str(sheet_hint or ""))
                       .replace("{text_layer}",
                                (instruction_page.get("text_layer") or "")
                                [:config.TEXT_LAYER_MAX_CHARS]))
        return call_json(
            system_prompt=TEXT_STRUCTURING_SYSTEM_PROMPT,
            user_text=instruction,
            images_b64=None,
            max_tokens=config.EXTRACT_MAX_TOKENS,
            model=config.AGENT_EXTRACT_MODEL,
            usage_tracker=self.usage,
            usage_stage="agent.extract_text",
            reasoning_effort=config.EXTRACT_REASONING_EFFORT or None,
            reasoning_max_tokens=config.EXTRACT_REASONING_MAX_TOKENS or None,
        )

    def run(self, scope: AgentScope) -> AgentResult:
        from backend.text_coverage import (fallback_objects, merge_objects,
                                           recover_sheet_number, text_coverage)
        try:
            page = scope.payload["page"]
            hint = scope.payload.get("sheet_hint") or ""
            page_text = page.get("text_layer")
            instruction = EXTRACTOR_USER_INSTRUCTION.replace(
                "{sheet_hint}", str(hint)) + _text_layer_block(page)

            # Rung 1: vision pass (unchanged behaviour, incl. compact retry)
            parsed = _wrap_bare_list(self._call(instruction, page),
                                     page["page_number"])
            if not isinstance(parsed, dict):
                # Second chance: same page, compact instructions. Runs only
                # when the full-schema pass returned nothing usable.
                print(f"[Extract] Page {page['page_number']}: full extraction "
                      f"failed, retrying compact")
                parsed = _wrap_bare_list(
                    self._call(instruction + _COMPACT_RETRY_SUFFIX, page),
                    page["page_number"],
                )
            if not isinstance(parsed, dict):
                # Don't give up on the page - the ladder below can still
                # rescue it from the text layer.
                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 keeps every object it found
            # (graphical_basis content exists only in the image); the text
            # pass fills in the text content the vision pass 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 = _wrap_bare_list(
                    self._text_structuring_call(page, hint), page["page_number"])
                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"])
                    # Fill header gaps the vision pass left null
                    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 - dark sheets are impossible.
            # Also merged (deduped) so stub notes never double up with
            # objects the earlier rungs already captured.
            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"]
                # _normalize_sheet only stamps its own "text_layer" rescue
                # grounding; restore the explicit fallback provenance.
                for stub in stubs:
                    stub["grounding"] = "text_layer_fallback"
                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 AgentResult(scope_id=scope.scope_id, artifacts=[sheet])
        except Exception as exc:
            return failure(scope, exc)


class SheetIndexAgent:
    name = "sheet_index"

    def __init__(self, usage: AgentUsage) -> None:
        self.usage = usage

    def run(self, scope: AgentScope) -> AgentResult:
        try:
            sheets = scope.payload.get("sheets") or []
            instruction = SHEET_INDEX_USER_INSTRUCTION.replace(
                "{sheet_index_input}",
                json.dumps(_index_input(sheets), ensure_ascii=True),
            )
            parsed = call_json(
                system_prompt=SHEET_INDEX_SYSTEM_PROMPT,
                user_text=instruction,
                max_tokens=config.SHEET_INDEX_MAX_TOKENS,
                model=config.AGENT_INDEX_MODEL,
                usage_tracker=self.usage,
                usage_stage="agent.sheet_index",
            )
            if isinstance(parsed, list):
                parsed = {"sheet_index": parsed, "missing_expected_sheets": []}
            if not isinstance(parsed, dict):
                raise ValueError("no sheet index returned")
            return AgentResult(scope_id=scope.scope_id, artifacts=[parsed])
        except Exception as exc:
            return failure(scope, exc)


class JurisdictionAgent:
    name = "jurisdiction"

    def __init__(self, usage: AgentUsage) -> None:
        self.usage = usage

    def run(self, scope: AgentScope) -> AgentResult:
        try:
            project_input: Dict = scope.payload.get("project_input") or {}
            instruction = JURISDICTION_USER_INSTRUCTION.replace(
                "{project_input}", json.dumps(project_input, ensure_ascii=True)
            )
            parsed = call_json(
                system_prompt=JURISDICTION_SYSTEM_PROMPT,
                user_text=instruction,
                max_tokens=config.JURISDICTION_MAX_TOKENS,
                model=config.AGENT_JURISDICTION_MODEL,
                usage_tracker=self.usage,
                usage_stage="agent.jurisdiction",
            )
            if not isinstance(parsed, dict):
                raise ValueError("no jurisdiction profile returned")
            profile = parsed.get("project_code_profile")
            artifact = profile if isinstance(profile, dict) else parsed
            return AgentResult(scope_id=scope.scope_id, artifacts=[artifact])
        except Exception as exc:
            return failure(scope, exc)