from __future__ import annotations

import time
import uuid
from collections import defaultdict
from dataclasses import dataclass, field, asdict
from typing import Optional

from paths import PROMPTS_DIR, read_text

OCR_SYSTEM_TPL = read_text(PROMPTS_DIR / "ocr_system.txt")
OCR_USER_ZERO_SHOT_TPL = read_text(PROMPTS_DIR / "ocr_user_zero_shot.txt")
OCR_FEW_SHOT_HEADER_TPL = read_text(PROMPTS_DIR / "ocr_user_few_shot_header.txt")
OCR_FEW_SHOT_TARGET_TPL = read_text(PROMPTS_DIR / "ocr_user_few_shot_target.txt")
EXPERT_SYSTEM_TPL = read_text(PROMPTS_DIR / "expert_system.txt")
EXPERT_USER_TPL = read_text(PROMPTS_DIR / "expert_user.txt")
JUDGE_SYSTEM_TPL = read_text(PROMPTS_DIR / "judge_system.txt")
JUDGE_USER_TPL = read_text(PROMPTS_DIR / "judge_user.txt")

def _safe_format(template: str, **values) -> str:
    """Format with missing keys treated as empty string."""
    return template.format_map(defaultdict(str, values))


def render_ocr_system(
    language: str, guidelines: str,
    override: Optional[str] = None, mode: str = "lines",
    json_template: Optional[str] = None,
) -> str:
    if override:
        tpl = override
    else:
        tpl = OCR_SYSTEM_TPL_BY_MODE.get(mode, OCR_SYSTEM_TPL)
    return _safe_format(
        tpl,
        language=language or "(unspecified)",
        guidelines=guidelines or "",
        json_template=(json_template or "").strip() or '{"lines": ["line 1", "line 2"]}',
    )


def render_ocr_user_target(few_shot: bool, override: Optional[str] = None) -> str:
    if few_shot:
        return _safe_format(OCR_FEW_SHOT_TARGET_TPL)
    tpl = override if override else OCR_USER_ZERO_SHOT_TPL
    return _safe_format(tpl)


def render_ocr_few_shot_header(n_examples: int) -> str:
    return _safe_format(OCR_FEW_SHOT_HEADER_TPL, n_examples=n_examples)


def render_expert_system(language: str, guidelines: str, override: Optional[str] = None) -> str:
    tpl = override if override else EXPERT_SYSTEM_TPL
    return _safe_format(tpl, language=language or "(unspecified)", guidelines=guidelines or "")


def render_expert_user(ocr_text: str, override: Optional[str] = None) -> str:
    tpl = override if override else EXPERT_USER_TPL
    return _safe_format(tpl, ocr_text=ocr_text)


def render_judge_system(language: str, guidelines: str, override: Optional[str] = None) -> str:
    tpl = override if override else JUDGE_SYSTEM_TPL
    return _safe_format(tpl, language=language or "(unspecified)", guidelines=guidelines or "")


def render_judge_user(
    *,
    ocr_text: str,
    expert_a_text: str,
    expert_a_conf: float,
    expert_a_corrections: list[str],
    expert_b_text: str,
    expert_b_conf: float,
    expert_b_corrections: list[str],
    override: Optional[str] = None,
) -> str:
    tpl = override if override else JUDGE_USER_TPL
    return _safe_format(
        tpl,
        ocr_text=ocr_text,
        expert_a_text=expert_a_text,
        expert_a_conf=f"{expert_a_conf:.2f}",
        expert_a_corrections="; ".join(expert_a_corrections) or "(none)",
        expert_b_text=expert_b_text,
        expert_b_conf=f"{expert_b_conf:.2f}",
        expert_b_corrections="; ".join(expert_b_corrections) or "(none)",
    )


@dataclass
class ICLExample:
    id: str
    image_b64: str
    text: str
    language: str = ""
    source: str = "corrected"
    added_at: float = field(default_factory=lambda: time.time())

    def to_public(self) -> dict:
        d = asdict(self)
        d["preview"] = (self.text[:120] + "…") if len(self.text) > 120 else self.text
        d["n_lines"] = self.text.count("\n") + 1 if self.text else 0
        return d

    def to_jsonl_dict(self) -> dict:
        return {
            "id": self.id,
            "image_b64": self.image_b64,
            "text": self.text,
            "language": self.language,
            "source": self.source,
            "added_at": self.added_at,
        }


class ICLPool:
    def __init__(self):
        self._items: list[ICLExample] = []

    def __len__(self) -> int:
        return len(self._items)

    @property
    def items(self) -> list[ICLExample]:
        return list(self._items)

    def add(self, *, image_b64: str, text: str, language: str = "", source: str = "corrected") -> ICLExample:
        for it in self._items:
            if it.image_b64 == image_b64:
                it.text = text
                it.language = language or it.language
                it.source = source
                it.added_at = time.time()
                return it
        ex = ICLExample(
            id=uuid.uuid4().hex[:12],
            image_b64=image_b64,
            text=text,
            language=language,
            source=source,
        )
        self._items.append(ex)
        return ex

    def remove(self, item_id: str) -> bool:
        before = len(self._items)
        self._items = [it for it in self._items if it.id != item_id]
        return len(self._items) < before

    def filter(self, language: str) -> list[ICLExample]:
        if not language:
            return list(self._items)
        return [it for it in self._items if not it.language or it.language == language]

    def sample(self, n: int, language: str = "") -> list[ICLExample]:
        if n <= 0:
            return []
        pool = self.filter(language)
        pool = sorted(pool, key=lambda it: it.added_at, reverse=True)
        return pool[:n]

    def to_jsonl_dicts(self) -> list[dict]:
        return [it.to_jsonl_dict() for it in self._items]

    def public_view(self) -> list[dict]:
        return [it.to_public() for it in self._items]


# ───────────────────────────────────────────────────────────────────────────
#  Prompt presets (selectable from the UI, then editable)
# ───────────────────────────────────────────────────────────────────────────

# ───────────────────────────────────────────────────────────────────────────
#  OCR output modes — only TWO modes, on purpose, to keep the UI simple:
#    1. "lines"        → fixed simple shape  {"lines": ["...", "..."]}
#    2. "custom_json"  → user pastes a JSON template; the model must fill it.
#  Any specialised shape (with_expansions, with_page_meta, catalogue entry…)
#  is just a JSON_TEMPLATE_PRESETS entry the user can load and edit.
# ───────────────────────────────────────────────────────────────────────────

OCR_SYSTEM_TPL_LINES = OCR_SYSTEM_TPL  # the file template, asks for {"lines": [str]}

OCR_SYSTEM_TPL_CUSTOM_JSON = """You are an expert palaeographer and Handwritten Text Recognition (HTR) annotator.

Target language / script: {language}

Transcription policy (apply this strictly to every string value you produce):
{guidelines}

Task: read the attached page image and produce a SINGLE JSON object that matches EXACTLY the following template — same keys, same nesting, same array structure. Fill the placeholder values with the actual content from the page; do not add, remove, or rename any field.

Required JSON template (your reply must use this exact shape):
{json_template}

Output rules:
- Reply with a single JSON object and nothing else (no markdown, no prose, no code fences).
- Preserve UTF-8 characters exactly as written on the page.
- Use null when a field genuinely has no value on the page (e.g. no title).
- Use "[…]" inside string values for unreadable portions.
- If a field is an array of strings, fill it with one entry per physical line / item, top to bottom.
"""

OCR_SYSTEM_TPL_BY_MODE = {
    "lines": OCR_SYSTEM_TPL_LINES,
    "custom_json": OCR_SYSTEM_TPL_CUSTOM_JSON,
}


# ───────────────────────────────────────────────────────────────────────────
#  JSON template presets — the user loads one as a starting point, then edits
#  the fields they care about. The model is asked to fill this exact shape.
# ───────────────────────────────────────────────────────────────────────────

JSON_TEMPLATE_PRESETS = {
    "simple_lines": (
        '{\n'
        '  "lines": ["line 1", "line 2", "line 3"]\n'
        '}'
    ),
    "page_with_metadata": (
        '{\n'
        '  "no_page": "12r",\n'
        '  "titre": "Liber primus",\n'
        '  "lines": ["line 1", "line 2", "line 3"]\n'
        '}'
    ),
    "with_expansions": (
        '{\n'
        '  "lines": [\n'
        '    {"raw": "Sancti Ioãnis euãgelium", "expanded": "Sancti Ioannis euangelium"},\n'
        '    {"raw": "...", "expanded": "..."}\n'
        '  ]\n'
        '}'
    ),
    "catalogue_entry": (
        '{\n'
        '  "folio": "12r",\n'
        '  "auteur": "Augustinus",\n'
        '  "titre": "De civitate dei",\n'
        '  "incipit": "...",\n'
        '  "explicit": "...",\n'
        '  "langue": "Latin",\n'
        '  "lines": ["..."]\n'
        '}'
    ),
    "structured_page_typed_lines": (
        '{\n'
        '  "page_meta": {\n'
        '    "main_language": "Latin",\n'
        '    "script": "caroline minuscule",\n'
        '    "estimated_period": "XIIᵉ s.",\n'
        '    "layout": "single-column"\n'
        '  },\n'
        '  "lines": [\n'
        '    {"n": 1, "text": "...", "type": "title"},\n'
        '    {"n": 2, "text": "...", "type": "body"}\n'
        '  ]\n'
        '}'
    ),
}


# ───────────────────────────────────────────────────────────────────────────
#  Guidelines presets — these populate the {guidelines} placeholder.
#  Picking one of these is the *primary* way for the user to define the
#  transcription policy. Templates above only enforce STRUCTURE, not POLICY.
# ───────────────────────────────────────────────────────────────────────────

GUIDELINES_PRESETS = {
    "preserve_as_is": (
        "Preserve the original orthography exactly as written; do not modernise spelling.\n"
        "Preserve line breaks: one transcribed line per physical line on the page.\n"
        "Keep abbreviation marks, ligatures, long-s (ſ), historical letters and special signs as written.\n"
        "Do not silently expand abbreviations.\n"
        "Preserve capitalisation and punctuation as on the source."
    ),
    "expand_abbreviations": (
        "Expand every abbreviation silently into its resolved form (e.g. Ioãnis → Ioannis, p̃r → pater).\n"
        "Preserve historical orthography otherwise (do not modernise spelling).\n"
        "Preserve line breaks and capitalisation.\n"
        "Use square brackets [ ] only when you supply letters that are not in the source."
    ),
    "modernise_spelling": (
        "Normalise spelling to modern conventions for {language} while keeping the meaning intact.\n"
        "Expand abbreviations silently.\n"
        "Modernise punctuation if necessary.\n"
        "Preserve line breaks of the original page."
    ),
    "strict_diplomatic": (
        "Strict diplomatic transcription: reproduce EVERYTHING you see, including capitalisation,\n"
        "punctuation marks, deleted/struck-through text (mark with ⟨…⟩), interlinear additions (mark with \\…/),\n"
        "abbreviation signs, ligatures, decorations.\n"
        "Use [...] for unreadable portions and (...) when you tentatively restore something."
    ),
}


# ───────────────────────────────────────────────────────────────────────────
#  System-prompt presets for the OCR call. The default for each mode is
#  picked automatically by render_ocr_system(); these presets give the user
#  alternative wordings to load and edit.
# ───────────────────────────────────────────────────────────────────────────

OCR_SYSTEM_PRESETS = {
    "default_plain_lines": OCR_SYSTEM_TPL_LINES,
    "default_custom_json": OCR_SYSTEM_TPL_CUSTOM_JSON,
    "minimal_lines": (
        "Transcribe the page image in {language}. Apply this policy:\n{guidelines}\n\n"
        "Reply with JSON only: {{\"lines\": [\"line 1\", \"line 2\", ...]}}"
    ),
    "minimal_custom_json": (
        "Transcribe the page image in {language}. Apply this policy:\n{guidelines}\n\n"
        "Reply ONLY with a JSON object matching this exact shape:\n{json_template}"
    ),
}

OCR_USER_PRESETS = {
    "default": OCR_USER_ZERO_SHOT_TPL,
    "verbose": (
        "Please transcribe the attached page image, line by line, top to bottom.\n"
        "Apply the policy and the JSON shape specified in your system prompt."
    ),
    "concise": "Transcribe line by line. Reply with the JSON specified in the system prompt.",
}

EXPERT_SYSTEM_PRESETS = {
    "default__balanced": EXPERT_SYSTEM_TPL,
    "strict__minimal_changes": (
        "You are a senior palaeographer reviewing an HTR / OCR prediction.\n\n"
        "Target language / script: {language}\n\n"
        "Transcription policy (apply strictly):\n{guidelines}\n\n"
        "Be CONSERVATIVE: only change a token if you are highly confident the image shows something different. "
        "Pay extra attention to: u/v, i/j, long-s (ſ), abbreviation marks (macrons, tildes, ⁊), punctuation, ligatures, capitalisation.\n\n"
        "Output format — reply with a single JSON object only:\n"
        "{{\"corrected_text\": \"line 1\\nline 2\\n...\",\n"
        " \"confidence\": 0.0,\n"
        " \"corrections\": [\"description of correction 1\", \"...\"]}}"
    ),
    "lenient__rewrite_freely": (
        "You are a palaeography expert correcting an HTR / OCR prediction.\n\n"
        "Target language / script: {language}\n\n"
        "Transcription policy (apply strictly):\n{guidelines}\n\n"
        "Be willing to depart freely from the OCR when the image clearly disagrees. "
        "Prefer a fully-rewritten line over a partial patch when the prediction is badly corrupted.\n\n"
        "Output format — reply with a single JSON object only:\n"
        "{{\"corrected_text\": \"...\\n...\", \"confidence\": 0.0, \"corrections\": [\"...\"]}}"
    ),
}

EXPERT_USER_PRESETS = {
    "default": EXPERT_USER_TPL,
    "concise": "HTR prediction:\n{ocr_text}\n\nVerify against the image. Reply with the JSON specified in your system prompt.",
}

JUDGE_SYSTEM_PRESETS = {
    "default": JUDGE_SYSTEM_TPL,
    "consensus_first": (
        "You are the judge in a two-expert HTR / OCR correction pipeline.\n\n"
        "Target language / script: {language}\n\n"
        "Transcription policy (apply strictly):\n{guidelines}\n\n"
        "Workflow: if Expert A and Expert B agree on a line, take that reading. "
        "Only adjudicate on lines where they differ, picking the reading most faithful to the image. "
        "You may synthesise or override entirely.\n\n"
        "Output format — reply with a single JSON object only:\n"
        "{{\"final_text\": \"...\\n...\", \"confidence\": 0.0,\n"
        " \"source\": \"expert_a\" | \"expert_b\" | \"synthesis\" | \"original\",\n"
        " \"rationale\": \"one short sentence\"}}"
    ),
}

JUDGE_USER_PRESETS = {
    "default": JUDGE_USER_TPL,
}


def list_preset_names() -> dict:
    """Public view of all preset names for the frontend."""
    return {
        "guidelines": list(GUIDELINES_PRESETS.keys()),
        "json_template": list(JSON_TEMPLATE_PRESETS.keys()),
        "ocr_system": list(OCR_SYSTEM_PRESETS.keys()),
        "ocr_user": list(OCR_USER_PRESETS.keys()),
        "expert_system": list(EXPERT_SYSTEM_PRESETS.keys()),
        "expert_user": list(EXPERT_USER_PRESETS.keys()),
        "judge_system": list(JUDGE_SYSTEM_PRESETS.keys()),
        "judge_user": list(JUDGE_USER_PRESETS.keys()),
    }


def get_preset(family: str, name: str) -> str:
    """Look up a single preset body. Returns empty string if not found."""
    table = {
        "guidelines": GUIDELINES_PRESETS,
        "json_template": JSON_TEMPLATE_PRESETS,
        "ocr_system": OCR_SYSTEM_PRESETS,
        "ocr_user": OCR_USER_PRESETS,
        "expert_system": EXPERT_SYSTEM_PRESETS,
        "expert_user": EXPERT_USER_PRESETS,
        "judge_system": JUDGE_SYSTEM_PRESETS,
        "judge_user": JUDGE_USER_PRESETS,
    }.get(family) or {}
    return table.get(name, "")