# tools/ocr_dispatch.py - OCR method selection (simple/complex/kraken) and post-process helpers.
# Runners use lazy imports from vision_ocr and ocr_backends to avoid circular imports.
from __future__ import annotations

import ast
import re

from tools._logging import _debug_log

_FORCE_SIMPLE_RE = re.compile(
    r"\b(use\s+(?:basic|simple)\s+ocr|force\s+simple|basic\s+ocr|simple\s+ocr|use\s+vision|vision\s+only|quick\s+ocr\s+only)\b",
    re.IGNORECASE,
)


def _is_pdf_path(file_path: str | None) -> bool:
    if not file_path or not isinstance(file_path, str):
        return False
    return file_path.strip().lower().endswith(".pdf")


def _is_low_quality_ocr_text(text: str) -> bool:
    """Heuristic quality gate used for simple->complex fallback."""
    if not isinstance(text, str):
        return True
    s = text.strip()
    if not s:
        return True
    if len(s) < 80:
        return True
    ratio = _letter_ratio(s)
    if ratio < 0.35:
        return True
    lines = [ln.strip() for ln in s.splitlines() if ln.strip()]
    if len(lines) >= 6:
        short_lines = sum(1 for ln in lines if len(ln) <= 2)
        if short_lines / len(lines) > 0.4:
            return True
    noise_chars = sum(1 for c in s if c in "|~•®°")
    if noise_chars / max(len(s), 1) > 0.02:
        return True
    return False


def _user_forced_simple(user_text: str) -> bool:
    return bool(_FORCE_SIMPLE_RE.search((user_text or "").strip()))


def _sanitize_ocr_output(result: str) -> str:
    """Unwrap wrapper-style OCR payloads like GenerationResult(text='...')."""
    if not isinstance(result, str):
        return result
    s = result.strip()
    if not s or "Result(" not in s:
        return result

    def _regex_unwrap_text(payload: str) -> str:
        wrapper = re.search(
            r"(?:^|\b)(?:[A-Za-z_][\w.]*\.)?(?:GenerationResult|OCRResult|Result)\s*\(",
            payload,
            re.DOTALL,
        )
        if not wrapper:
            return result
        body = payload[wrapper.end() :]
        m = re.search(r"\btext\s*=\s*(['\"])((?:\\.|(?!\1).)*)\1", body, re.DOTALL)
        if not m:
            return result
        q = m.group(1)
        raw = m.group(2)
        try:
            # Reuse Python string parser semantics so escaped newlines/quotes are restored.
            return ast.literal_eval(f"{q}{raw}{q}")
        except Exception:
            return raw

    try:
        node = ast.parse(s, mode="eval").body
    except Exception:
        return _regex_unwrap_text(s)
    if not isinstance(node, ast.Call):
        return _regex_unwrap_text(s)
    fn_name = ""
    if isinstance(node.func, ast.Name):
        fn_name = node.func.id
    elif isinstance(node.func, ast.Attribute):
        fn_name = node.func.attr
    if fn_name not in ("GenerationResult", "OCRResult", "Result"):
        return _regex_unwrap_text(s)
    for kw in node.keywords or []:
        if kw.arg == "text" and isinstance(kw.value, ast.Constant) and isinstance(kw.value.value, str):
            return kw.value.value
    return _regex_unwrap_text(s)


def _letter_ratio(s: str) -> float:
    """Fraction of characters that are letters (a-z A-Z)."""
    if not s:
        return 0.0
    letters = sum(1 for c in s if c.isalpha())
    return letters / len(s)


def _footnote_sort_key(key: str):
    """Sort footnote keys: numbers first (1,2,...), then symbols (*, †, ‡, §, ¶)."""
    if isinstance(key, str) and key.isdigit():
        try:
            return (0, int(key))
        except ValueError:
            pass
    return (1, key if isinstance(key, str) else str(key))


def post_process_to_markdown(text: str, footnotes: dict) -> str:
    """Final cleanup: H1 title, paragraph joining, conservative bold, **Footnotes**."""
    text = re.sub(r"\n{3,}", "\n\n", text)
    text = text.strip()
    lines = text.split("\n")
    result = []
    i = 0
    while i < len(lines):
        line = lines[i]
        s = line.strip()
        if not s:
            result.append(line)
            i += 1
            continue
        while i + 1 < len(lines):
            next_s = lines[i + 1].strip()
            if not next_s:
                break
            if len(s) < 40 and not re.search(r"[.!?\"']\s*$", s):
                s = s + " " + next_s
                i += 1
            else:
                break
        result.append(s)
        i += 1
    text = "\n".join(result)
    first_line = text.split("\n")[0].strip()
    m = re.match(r"^([A-Z][A-Z\s]{2,60}?)(?=\s+[A-Z][a-z]|\s*$)", first_line)
    if m:
        title = m.group(1).strip()
        if title.isupper() and title.replace(" ", "").isalpha() and 5 <= len(title) <= 80:
            rest = text[len(first_line):].lstrip()
            text = f"# {title}\n\n{rest}"
    else:
        if 5 <= len(first_line) <= 60 and first_line.isupper() and first_line.replace(" ", "").isalpha():
            rest = "\n".join(text.split("\n")[1:]).lstrip()
            text = f"# {first_line}\n\n{rest}"
    text = re.sub(r"[\s_°:;]+$", "", text)

    def _format_heading(line: str, is_first: bool) -> str:
        s = line.strip()
        if 5 <= len(s) <= 80 and s.isupper() and s.replace(" ", "").isalpha():
            return f"# {s}" if is_first else f"**{s}**"
        return line

    lines = text.split("\n")
    result = []
    for i, line in enumerate(lines):
        prev_blank = i == 0 or not lines[i - 1].strip()
        next_blank = i == len(lines) - 1 or not lines[i + 1].strip()
        if prev_blank and next_blank and line.strip():
            result.append(_format_heading(line, is_first=(i == 0)))
        else:
            result.append(line)
    text = "\n".join(result)
    if footnotes:
        text += "\n\n---\n\n**Footnotes**\n\n"
        for key in sorted(footnotes.keys(), key=_footnote_sort_key):
            entries = footnotes[key]
            body = " ".join(re.sub(r"^Page\s+\d+:\s*", "", str(e).strip()) for e in entries if e)
            if not body:
                body = " ".join(str(e) for e in entries)
            body = re.sub(r"\s*&\s*=.*$", "", body)
            body = re.sub(r"\d+°\s*:?\s*$", "", body)
            body = re.sub(r"\s+:\s*$", "", body)
            body_stripped = body.strip()
            if body_stripped and body_stripped.replace(str(key), "").replace(" ", ""):
                text += f"[^{key}]: {body_stripped}\n"
    return text


def post_process_to_latex(text: str, footnotes: dict) -> str:
    """Basic Markdown → LaTeX conversion."""
    # Protect footnote refs [^...] so * replacement does not corrupt [^*]
    _placeholders: list[str] = []
    def _save_ref(m: re.Match) -> str:
        _placeholders.append(m.group(0))
        return f"\x00FN\x00{len(_placeholders) - 1}\x00FN\x00"
    text = re.sub(r"\[\^[^\]]+\]", _save_ref, text)
    text = text.replace("**", "\\textbf{").replace("*", "\\textit{")
    text = re.sub(r'\\textbf\{(.*?)\}', r'\\textbf{\1}', text)
    text = re.sub(r'\\textit\{(.*?)\}', r'\\textit{\1}', text)
    for i, ref in enumerate(_placeholders):
        text = text.replace(f"\x00FN\x00{i}\x00FN\x00", ref)
    if footnotes:
        text += "\n\n\\section*{Footnotes}\n"
        for key in sorted(footnotes.keys(), key=_footnote_sort_key):
            entries = footnotes[key]
            body = " ".join(re.sub(r"^Page\s+\d+:\s*", "", str(e).strip()) for e in entries if e)
            text += f"\\footnote[{key}]{{{body}}}\n"
    return text


def final_ocr_polish(text: str) -> str:
    """Light, broad cleanup pass — safe for most English printed documents. Fixes typography and spacing."""
    if not text or not text.strip():
        return text
    text = text.replace("\u201c", '"').replace("\u201d", '"')
    text = text.replace("\u2018", "'").replace("\u2019", "'")
    text = text.replace("``", '"').replace("''", '"')
    text = text.replace("\u2014", "--").replace("\u2013", "-")
    text = re.sub(r"([a-z])([.,;:!?])([A-Z])", r"\1\2 \3", text)

    def keep_line(line: str) -> bool:
        s = line.strip()
        if not s:
            return True
        if len(s) <= 2 and not s.isalpha():
            return False
        if len(s) <= 8 and _letter_ratio(s) < 0.2:
            return False
        return True

    text = "\n".join(line if keep_line(line) else "" for line in text.split("\n"))
    text = re.sub(r"\n{3,}", "\n\n", text)
    return text.strip()


# ----- Runners: lazy import from vision_ocr and ocr_backends -----

def _ocr_runner_simple():
    """Return a callable (file_path) -> result that uses ocr_scanned_pdf."""
    from tools.vision_ocr import ocr_scanned_pdf
    def run(file_path: str):
        if hasattr(ocr_scanned_pdf, "invoke"):
            return ocr_scanned_pdf.invoke({"file_path": file_path})
        return ocr_scanned_pdf(file_path)
    return run


def _ocr_runner_complex():
    """Return a callable (file_path) -> result that uses ocr_complex_document with defaults."""
    from tools.ocr_backends.tesseract import ocr_complex_document
    def run(file_path: str):
        kwargs = {
            "file_path": file_path,
            "output_format": "markdown",
            "enhance_for_tables": True,
            "detect_formatting": True,
            "detect_footnotes": True,
        }
        if hasattr(ocr_complex_document, "invoke"):
            return ocr_complex_document.invoke(kwargs)
        return ocr_complex_document(**kwargs)
    return run


def _ocr_runner_kraken():
    """Return a callable (file_path) -> result that uses ocr_with_kraken_api with defaults."""
    from tools.ocr_backends.kraken import ocr_with_kraken_api
    def run(file_path: str):
        if hasattr(ocr_with_kraken_api, "invoke"):
            return ocr_with_kraken_api.invoke({
                "file_path": file_path,
                "output_format": "markdown",
                "model_path": None,
            })
        return ocr_with_kraken_api(file_path=file_path, output_format="markdown", model_path=None)
    return run


def choose_ocr_method(
    user_text: str,
    file_path: str | None = None,
) -> tuple:
    """
    Choose OCR function (ocr_scanned_pdf vs ocr_complex_document vs ocr_with_kraken_api) from user message.
    Returns (function_to_call, reason_explanation, method_name).
    """
    from tools.ocr_backends.kraken import KRAKEN_AVAILABLE as _KRAKEN_AVAILABLE
    text = (user_text or "").lower().strip()
    _debug_log(
        f"[choose_ocr_method] user_text={repr((user_text or '')[:200])}, normalized_len={len(text)}, _KRAKEN_AVAILABLE={_KRAKEN_AVAILABLE}",
        "INFO",
    )
    if not text:
        _debug_log("[choose_ocr_method] -> complex (no user message)", "INFO")
        return _ocr_runner_complex(), "No user message; using complex OCR (default).", "complex"

    kraken_explicit = re.search(r"\b(use\s+kraken|kraken\s+ocr|kraken\s+api|run\s+kraken)\b", text, re.I)
    if _KRAKEN_AVAILABLE and kraken_explicit:
        _debug_log("[choose_ocr_method] -> kraken (user asked for Kraken OCR)", "INFO")
        return _ocr_runner_kraken(), "User asked for Kraken OCR.", "kraken"

    force_complex = re.search(
        r"\b(use\s+complex\s+ocr|force\s+complex|complex\s+ocr|high\s+quality\s+ocr|use\s+tesseract|tesseract\s+ocr)\b",
        text, re.IGNORECASE,
    )
    force_simple = _FORCE_SIMPLE_RE.search(text)
    if force_complex:
        return _ocr_runner_complex(), "User asked for complex/high-quality OCR.", "complex"
    if force_simple:
        return _ocr_runner_simple(), "User asked for basic/simple OCR.", "simple"

    negation_simple = re.search(
        r"\b(no\s+tables?|no\s+formatting|no\s+layout|just\s+plain\s+text|plain\s+text\s+only|"
        r"simple\s+scan|basic\s+scan|text\s+only|no\s+need\s+for\s+formatting|"
        r"don'?t\s+need\s+formatting|without\s+formatting|no\s+footnotes?|no\s+charts?|no\s+graphs?|extract\s+text\s+only)\b",
        text, re.IGNORECASE,
    )
    if negation_simple:
        return _ocr_runner_simple(), "User indicated plain/simple extraction (no tables/formatting).", "simple"

    simple_preference = re.search(
        r"\b(quick(?:ly)?|fast|just\s+text|extract\s+text\s+only|no\s+formatting\s+needed|minimal\s+output|raw\s+text|text\s+extraction\s+only)\b",
        text, re.IGNORECASE,
    )
    if simple_preference:
        return _ocr_runner_simple(), "User asked for quick/fast or text-only extraction.", "simple"

    quality_complex = re.search(
        r"\b(improve|better\s+ocr|more\s+accurate|preserve\s+(?:style|formatting|layout)|"
        r"with\s+formatting|keep\s+layout|high\s+quality|best\s+quality|accurate\s+ocr|correct\s+formatting|fix\s+formatting)\b",
        text, re.IGNORECASE,
    )
    if quality_complex:
        return _ocr_runner_complex(), "User asked for better/accurate OCR or to preserve formatting.", "complex"

    if _KRAKEN_AVAILABLE and re.search(
        r"\b(old\s+book|archaic|historical\s+document|antique\s+book|early\s+print|glyphs?|long\s+s|medieval|manuscript|high\s+performance\s+ocr|best\s+accuracy)\b",
        text, re.IGNORECASE,
    ):
        return _ocr_runner_kraken(), "User mentioned old book / historical / archaic document; using Kraken.", "kraken"

    content_complex = re.search(
        r"\b(tables?|charts?|graphs?|footnotes?|bold|italic|formatted|formatting|layout|"
        r"columns?|rows?|tabular|spreadsheet(?:\s*like)?|with\s+figures?|includes?\s+diagrams?|has\s+references?|"
        r"academic\s+paper|research\s+pdf|invoices?\s+with\s+columns?|extract\s+table|preserve\s+formatting|"
        r"keep\s+layout|legal\s+document|financial\s+document|printed\s+book|scanned\s+book|religious\s+text|academic\s+text|book\s+with\s+footnotes?)\b",
        text, re.IGNORECASE,
    )
    if content_complex:
        return _ocr_runner_complex(), "User mentioned tables/formatting/layout or complex document type.", "complex"

    if re.search(r"\b(correct|fix|improve|accurate|proper|better)\b", text, re.I) and re.search(r"\b(ocr|scan|pdf|document|text)\b", text, re.I):
        return _ocr_runner_complex(), "User asked to correct/improve OCR or document quality.", "complex"
    if re.search(r"\b(with\s+formatting|with\s+layout|with\s+tables?|with\s+footnotes?)\b", text, re.I):
        return _ocr_runner_complex(), "User asked for formatting/layout/tables/footnotes.", "complex"

    _debug_log("[choose_ocr_method] -> complex (default, no strong signal)", "INFO")
    return _ocr_runner_complex(), "No strong signal; using complex OCR (default for best quality).", "complex"


def choose_ocr_method_with_confidence(user_text: str, file_path: str | None = None) -> tuple:
    """Like choose_ocr_method but also returns a confidence score in [0, 1]."""
    from tools.ocr_backends.kraken import KRAKEN_AVAILABLE as _KRAKEN_AVAILABLE
    text = (user_text or "").lower().strip()
    if not text:
        return _ocr_runner_complex(), "No user message; using complex OCR (default).", "complex", 0.6
    if _KRAKEN_AVAILABLE and re.search(r"\b(use\s+kraken|kraken\s+ocr|kraken\s+api)\b", text, re.I):
        return _ocr_runner_kraken(), "User asked for Kraken OCR.", "kraken", 1.0
    if re.search(r"\b(use\s+complex\s+ocr|force\s+complex|complex\s+ocr|high\s+quality\s+ocr|use\s+tesseract|tesseract\s+ocr)\b", text, re.I):
        return _ocr_runner_complex(), "User asked for complex/high-quality OCR.", "complex", 1.0
    if re.search(r"\b(use\s+(?:basic|simple)\s+ocr|force\s+simple|basic\s+ocr|simple\s+ocr|use\s+vision|vision\s+only)\b", text, re.I):
        return _ocr_runner_simple(), "User asked for basic/simple OCR.", "simple", 1.0
    if re.search(r"\b(no\s+tables?|no\s+formatting|just\s+plain\s+text|simple\s+scan|text\s+only|without\s+formatting)\b", text, re.I):
        return _ocr_runner_simple(), "User indicated plain/simple extraction.", "simple", 0.95
    if re.search(r"\b(quick(?:ly)?|fast|just\s+text|no\s+formatting\s+needed)\b", text, re.I):
        return _ocr_runner_simple(), "User asked for quick/fast or text-only.", "simple", 0.9
    if _KRAKEN_AVAILABLE and re.search(r"\b(old\s+book|archaic|historical\s+document|antique|medieval|manuscript|glyphs?)\b", text, re.I):
        return _ocr_runner_kraken(), "User mentioned old/historical document; using Kraken.", "kraken", 0.85
    if re.search(r"\b(improve|better\s+ocr|preserve\s+(?:style|formatting|layout)|with\s+formatting|keep\s+layout|high\s+quality)\b", text, re.I):
        return _ocr_runner_complex(), "User asked for better/accurate OCR or preserve formatting.", "complex", 0.95
    if re.search(r"\b(tables?|charts?|footnotes?|bold|italic|formatted|layout|columns?|tabular|spreadsheet|academic\s+paper|extract\s+table)\b", text, re.I):
        return _ocr_runner_complex(), "User mentioned tables/formatting or complex document.", "complex", 0.85
    return _ocr_runner_complex(), "No strong signal; using complex OCR (default for best quality).", "complex", 0.6


def run_ocr_with_dispatch(file_path: str, user_text: str = "", log_reason: bool = True) -> str:
    """Run OCR using the method chosen from user_text. Use when you have the user's message and want to auto-select simple vs complex OCR."""
    fn, reason, method_name = choose_ocr_method(user_text, file_path)
    if log_reason:
        _debug_log(f"[OCR dispatch] {method_name}: {reason}", "DEBUG")
    result = _sanitize_ocr_output(fn(file_path))
    if (
        method_name == "simple"
        and _is_pdf_path(file_path)
        and not _user_forced_simple(user_text)
        and _is_low_quality_ocr_text(result)
    ):
        if log_reason:
            _debug_log(
                "[OCR dispatch] simple OCR output looked low-quality; retrying with complex OCR.",
                "INFO",
            )
        complex_runner = _ocr_runner_complex()
        complex_result = _sanitize_ocr_output(complex_runner(file_path))
        if isinstance(complex_result, str) and complex_result.strip():
            return complex_result
    return result


def _test_choose_ocr_method():
    """Test cases for choose_ocr_method: (user_text, expected_method, reason_substring). Run with: python -m tools.ocr_dispatch"""
    cases = [
        ("Use complex OCR on this PDF", "complex", "complex"),
        ("Force simple OCR please", "simple", "simple"),
        ("Use basic ocr", "simple", "basic"),
        ("OCR this but no tables, just plain text", "simple", "plain"),
        ("Simple scan only", "simple", "simple"),
        ("Extract text only, no formatting needed", "simple", "text-only"),
        ("Quick OCR this file", "simple", "quick"),
        ("I need fast extraction", "simple", "fast"),
        ("Improve the OCR and preserve formatting", "complex", "preserve"),
        ("Better OCR with layout", "complex", "layout"),
        ("I need accurate OCR and keep layout", "complex", "accurate"),
        ("This PDF has tables and footnotes", "complex", "tables"),
        ("OCR the invoice with columns", "complex", "columns"),
        ("Academic paper with references", "complex", "academic"),
        ("Extract the table from this scan", "complex", "table"),
        ("Document has bold and italic text", "complex", "bold"),
        ("", "complex", "default"),
        ("OCR this file", "complex", "default"),
    ]
    failed = []
    for user_text, expected_method, reason_substring in cases:
        fn, reason, method_name = choose_ocr_method(user_text)
        if method_name != expected_method:
            failed.append((user_text, f"expected {expected_method}, got {method_name}"))
        elif reason_substring != "default" and reason_substring not in reason.lower():
            failed.append((user_text, f"reason should contain {reason_substring!r}: {reason}"))
        assert callable(fn), "function_to_call must be callable"
    if failed:
        for msg, err in failed:
            print(f"FAIL: {msg!r} -> {err}")
        raise AssertionError(f"{len(failed)} test(s) failed")
    print("choose_ocr_method: all tests passed.")


__all__ = [
    "_letter_ratio",
    "_footnote_sort_key",
    "post_process_to_markdown",
    "post_process_to_latex",
    "final_ocr_polish",
    "_ocr_runner_simple",
    "_ocr_runner_complex",
    "_ocr_runner_kraken",
    "choose_ocr_method",
    "choose_ocr_method_with_confidence",
    "run_ocr_with_dispatch",
    "_test_choose_ocr_method",
]


if __name__ == "__main__":
    _test_choose_ocr_method()
