"""Text extraction from PDF, DOCX, DOC, and RTF files.

Primary PDF extraction: PyMuPDF direct text (for born-digital PDFs).
Fallback: Mistral OCR (for scanned documents or broken OCR layers).

Routing logic (document-level, not per-page):
  1. PyMuPDF extracts text from every page.
  2. Pages are quality-checked via _text_quality_ok().
  3. If ALL pages pass → use PyMuPDF output (free, ~50ms, no API call).
  4. If ANY page fails → call Mistral OCR once for the entire PDF.
     Mistral returns per-page Markdown; page_offsets are computed from it.

DOC files: converted to DOCX via LibreOffice before extraction.
Post-processing: Hebrew abbreviation quote fixer (PyMuPDF path only;
Mistral handles gershayim natively).
"""

from __future__ import annotations

import asyncio
import base64
import io
import logging
import re
import subprocess
import tempfile
from pathlib import Path

import fitz  # PyMuPDF
import httpx
from PIL import Image
from docx import Document as DocxDocument
from striprtf.striprtf import rtf_to_text

from legal_mcp import config
from legal_mcp.services import storage

logger = logging.getLogger(__name__)

# ── Mistral OCR ───────────────────────────────────────────────────

_MISTRAL_OCR_URL = "https://api.mistral.ai/v1/ocr"
_MISTRAL_OCR_MODEL = "mistral-ocr-latest"


async def _call_mistral_ocr(path: Path) -> list[str]:
    """Call Mistral OCR API on a PDF. Returns per-page Markdown text list.

    The Mistral response contains ``pages[i].markdown`` for each page.
    If the response has fewer pages than the PDF, trailing pages are
    padded with empty strings by the caller.
    """
    if not config.MISTRAL_API_KEY:
        raise RuntimeError(
            "MISTRAL_API_KEY not configured — cannot OCR scanned PDF. "
            "Set the env var in Coolify."
        )

    pdf_b64 = base64.b64encode(path.read_bytes()).decode()

    async with httpx.AsyncClient(timeout=300.0) as client:
        resp = await client.post(
            _MISTRAL_OCR_URL,
            headers={
                "Authorization": f"Bearer {config.MISTRAL_API_KEY}",
                "Content-Type": "application/json",
            },
            json={
                "model": _MISTRAL_OCR_MODEL,
                "document": {
                    "type": "document_url",
                    "document_url": f"data:application/pdf;base64,{pdf_b64}",
                },
                "include_image_base64": False,
            },
        )

    if resp.status_code != 200:
        raise RuntimeError(
            f"Mistral OCR returned {resp.status_code}: {resp.text[:400]}"
        )

    pages = resp.json().get("pages", [])
    return [p.get("markdown", "") for p in pages]


# ── Hebrew text quality detection ────────────────────────────────

_HEBREW_RE = re.compile(r'[֐-׿]')
_WORD_RE = re.compile(r'\S+')


def _text_quality_ok(text: str) -> bool:
    """Check if PyMuPDF-extracted text is genuine Hebrew legal content.

    Returns True if text appears to be real content.
    Broken OCR layers from scanned PDFs often have:
    - Very short words / single-character fragments
    - Each word on its own line (high words-per-line ratio)
    - Non-Hebrew characters mixed in
    """
    words = _WORD_RE.findall(text)
    if len(words) < 10:
        return False

    avg_len = sum(len(w) for w in words) / len(words)
    if avg_len < 2.5:
        return False

    single_char_pct = sum(1 for w in words if len(w) == 1) / len(words)
    if single_char_pct > 0.4:
        return False

    lines = [l for l in text.split("\n") if l.strip()]
    if lines and len(words) / len(lines) < 3.0:
        return False

    letters = re.findall(r'[a-zA-Z֐-׿]', text)
    if letters:
        hebrew_pct = sum(1 for c in letters if _HEBREW_RE.match(c)) / len(letters)
        if hebrew_pct < 0.5:
            return False

    return True


# ── Hebrew abbreviation quote fixer (PyMuPDF path only) ──────────

_HEBREW_ABBREV_FIXES: dict[str, str] = {
    'עוהייד': 'עוה"ד',
    'עוייד': 'עו"ד',
    'הנייל': 'הנ"ל',
    'מצייב': 'מצ"ב',
    'ביהמייש': 'ביהמ"ש',
    'תייז': 'ת"ז',
    'עייי': 'ע"י',
    'אחייכ': 'אח"כ',
    'סייק': 'ס"ק',
    'דייר': 'ד"ר',
    'כדוייח': 'כדו"ח',
    'חווייד': 'חוו"ד',
    'מייר': 'מ"ר',
    'יחייד': 'יח"ד',
    'בייכ': 'ב"כ',
    # Patterns where double-yod (יי) substitutes for gershayim (״) in born-digital PDFs
    'בליימ': 'בל"מ',
    'תמייא': 'תמ"א',
}

_ABBREV_PATTERN = re.compile(
    '|'.join(re.escape(k) for k in sorted(_HEBREW_ABBREV_FIXES, key=len, reverse=True))
)

_HEBREW_YEAR_RE = re.compile(r'(תש[א-ת]+)יי([א-ת])')


def _fix_hebrew_quotes(text: str) -> str:
    """Fix gershayim encoded as double-yod in born-digital PDFs."""
    text = _ABBREV_PATTERN.sub(lambda m: _HEBREW_ABBREV_FIXES[m.group()], text)
    text = _HEBREW_YEAR_RE.sub(r'\1"\2', text)
    return text


# ── Page joining ──────────────────────────────────────────────────

# Separator used when joining per-page text. Constant so chunker /
# retrofit can reproduce the join when computing page offsets.
PAGE_SEPARATOR = "\n\n"


def _join_pages(pages_text: list[str]) -> tuple[str, list[int]]:
    """Join per-page text with PAGE_SEPARATOR while recording start offsets."""
    offsets: list[int] = []
    parts: list[str] = []
    cursor = 0
    for i, pg in enumerate(pages_text):
        offsets.append(cursor)
        parts.append(pg)
        cursor += len(pg)
        if i < len(pages_text) - 1:
            parts.append(PAGE_SEPARATOR)
            cursor += len(PAGE_SEPARATOR)
    return "".join(parts), offsets


# ── PDF extraction ────────────────────────────────────────────────


async def _extract_pdf(path: Path) -> tuple[str, int, list[int]]:
    """Extract text from PDF using document-level routing.

    Stage 1 — PyMuPDF pre-screen (free, ~50ms, no API call):
        Run on every page. Collect per-page text; flag pages where
        PyMuPDF returns < 50 chars or _text_quality_ok() fails
        (scanned, blank, or broken embedded OCR layer).

    Stage 2 — Mistral OCR (triggered when any page fails Stage 1):
        Send the entire PDF to Mistral once. Use its per-page Markdown
        for ALL pages — consistent source, no mixed plain/Markdown formats.
        Mistral handles gershayim natively; no quote-fix applied.

    Page offsets are always computed so the chunker can attribute each
    chunk to its source page number (multimodal hybrid retrieval).
    """
    doc = fitz.open(str(path))
    page_count = len(doc)

    # Stage 1: PyMuPDF pre-screen
    pymupdf_pages: list[str] = []
    failed: list[int] = []

    for i in range(page_count):
        text = doc[i].get_text().strip()
        if len(text) > 50 and _text_quality_ok(text):
            pymupdf_pages.append(_fix_hebrew_quotes(text))
        else:
            pymupdf_pages.append("")
            failed.append(i)

    doc.close()

    if not failed:
        logger.debug(
            "PDF %s: all %d pages digital — PyMuPDF only", path.name, page_count
        )
        joined, offsets = _join_pages(pymupdf_pages)
        return joined, page_count, offsets

    # Stage 2: Mistral OCR for entire document
    logger.info(
        "PDF %s: %d/%d pages failed quality check → Mistral OCR",
        path.name, len(failed), page_count,
    )
    mistral_pages = await _call_mistral_ocr(path)

    # Pad if Mistral returns fewer pages than PyMuPDF counted
    while len(mistral_pages) < page_count:
        mistral_pages.append("")

    joined, offsets = _join_pages(mistral_pages[:page_count])
    return joined, page_count, offsets


def page_at_offset(offset: int, page_offsets: list[int]) -> int:
    """Return the 1-based page number containing a given char offset.

    page_offsets[i] is the start of page (i+1) in the joined text.
    """
    if not page_offsets:
        return 1
    page = 1
    for i, start in enumerate(page_offsets):
        if start <= offset:
            page = i + 1
        else:
            break
    return page


# ── Public entry point ────────────────────────────────────────────


async def extract_text(file_path: str) -> tuple[str, int, list[int] | None]:
    """Extract text from a document file.

    Returns:
        ``(text, page_count, page_offsets)`` where:
        - ``text``: extracted text. Plain text for PyMuPDF path;
          Markdown for Mistral path (tables, ``##`` headers preserved).
        - ``page_count``: number of pages (0 for non-PDF)
        - ``page_offsets``: char start of each page inside ``text``,
          or ``None`` for non-PDF formats
    """
    path = Path(file_path)
    suffix = path.suffix.lower()

    if suffix == ".pdf":
        return await _extract_pdf(path)
    elif suffix == ".docx":
        return _extract_docx(path), 0, None
    elif suffix == ".doc":
        return _extract_doc(path), 0, None
    elif suffix == ".rtf":
        return _extract_rtf(path), 0, None
    elif suffix in (".txt", ".md"):
        return path.read_text(encoding="utf-8"), 0, None
    else:
        raise ValueError(f"Unsupported file type: {suffix}")


# ── Non-PDF formats ───────────────────────────────────────────────


def _extract_doc(path: Path) -> str:
    """Extract text from legacy .doc via LibreOffice → DOCX conversion."""
    with tempfile.TemporaryDirectory() as tmp_dir:
        # Isolate the LibreOffice user profile per call: headless soffice
        # locks a single shared profile, so concurrent .doc conversions would
        # otherwise fail with a profile-lock error.
        result = subprocess.run(
            [
                "libreoffice",
                f"-env:UserInstallation=file://{tmp_dir}/lo-profile",
                "--headless", "--convert-to", "docx", str(path), "--outdir", tmp_dir,
            ],
            capture_output=True, text=True, timeout=120,
        )
        if result.returncode != 0:
            raise RuntimeError(f"LibreOffice conversion failed: {result.stderr}")
        docx_path = Path(tmp_dir) / f"{path.stem}.docx"
        if not docx_path.exists():
            raise FileNotFoundError(f"Converted file not found: {docx_path}")
        return _extract_docx(docx_path)


def _extract_docx(path: Path) -> str:
    """Extract text from DOCX file."""
    doc = DocxDocument(str(path))
    paragraphs = [p.text for p in doc.paragraphs if p.text.strip()]
    return "\n\n".join(paragraphs)


def _extract_rtf(path: Path) -> str:
    """Extract text from RTF file."""
    rtf_content = path.read_text(encoding="utf-8", errors="replace")
    return rtf_to_text(rtf_content)


# ── Multimodal page rendering (V9) ───────────────────────────────
# Unchanged — multimodal embedding always uses PyMuPDF-rendered images
# regardless of whether text extraction used PyMuPDF or Mistral.


def _pixmap_to_pil(pix: fitz.Pixmap) -> Image.Image:
    """Convert a PyMuPDF pixmap to PIL.Image (RGB)."""
    if pix.alpha:
        pix = fitz.Pixmap(pix, 0)
    return Image.frombytes("RGB", (pix.width, pix.height), pix.samples)


def render_pages_for_multimodal(
    pdf_path: str | Path,
    embed_dpi: int,
    thumb_dpi: int | None = None,
    thumbnail_dir: Path | None = None,
) -> list[tuple[Image.Image, Path | None]]:
    """Render each PDF page as PIL.Image at ``embed_dpi`` for the
    multimodal embedder, and optionally save JPEG thumbnails.

    Returns ``[(pil_image, thumb_path_or_None), ...]`` in page order.
    """
    src = Path(pdf_path)
    if not src.is_file():
        raise FileNotFoundError(f"PDF not found: {src}")
    if thumbnail_dir is not None:
        thumbnail_dir.mkdir(parents=True, exist_ok=True)

    out: list[tuple[Image.Image, Path | None]] = []
    doc = fitz.open(str(src))
    try:
        for page_idx, page in enumerate(doc):
            page_num = page_idx + 1
            pix = page.get_pixmap(dpi=embed_dpi)
            img = _pixmap_to_pil(pix)

            thumb_path: Path | None = None
            if thumbnail_dir is not None and thumb_dpi:
                thumb_path = thumbnail_dir / f"p{page_num:03d}.jpg"
                ratio = thumb_dpi / embed_dpi
                thumb_size = (
                    max(1, int(img.width * ratio)),
                    max(1, int(img.height * ratio)),
                )
                thumb = img.resize(thumb_size, Image.Resampling.LANCZOS)
                _tbuf = io.BytesIO()
                thumb.save(_tbuf, "JPEG", quality=75, optimize=True)
                try:
                    _tkey = thumb_path.resolve().relative_to(
                        Path(config.DATA_DIR).resolve()).as_posix()
                    storage.put_bytes_sync(
                        _tkey, _tbuf.getvalue(), bucket=storage.Bucket.DERIVED,
                        content_type="image/jpeg")
                except ValueError:
                    thumb.save(thumb_path, "JPEG", quality=75, optimize=True)

            out.append((img, thumb_path))
    finally:
        doc.close()
    return out


# ── Nevo preamble stripping ───────────────────────────────────────

_NEVO_MARKERS = ("ספרות:", "חקיקה שאוזכרה:", "מיני-רציו:", "פסקי דין שאוזכרו:",
                 "כתבי עת:", "הועתק מנבו")

_DECISION_START = re.compile(
    r"^[ \t>*_#]{0,6}(?:"
    r"בפנינו|לפנינו|לפניי|הערר שבנדון|ועדת הערר לתכנון|רקע עובדתי|עסקינן|"
    r"פסק[ \t\-]{0,3}די(?:ן|נו)|"
    r"(?:כב(?:וד)?['׳\"]?\s*)?(?:ה?שופט[ת]?|ה?נשיא[ה]?|המשנה לנשיא)\s+[^\n,]{1,40}:"
    r")",
    re.MULTILINE,
)


def strip_nevo_preamble(text: str) -> str:
    """Remove Nevo database preamble (bibliography, legislation, mini-ratio).

    Returns the original text unchanged if no preamble is detected.
    Works on both plain text (PyMuPDF) and Markdown (Mistral) since
    _DECISION_START already tolerates leading ``[ \t>*_#]{0,6}``.
    """
    head = text[:1500]
    if not any(marker in head for marker in _NEVO_MARKERS):
        return text
    m = _DECISION_START.search(text)
    if m and m.start() > 50:
        stripped = text[m.start():]
        logger.debug("Stripped %d chars of Nevo preamble", m.start())
        return stripped
    return text


_RATIO_MARKER = "מיני-רציו:"


def extract_nevo_ratio(text: str) -> str:
    """Return the Nevo מיני-רציו block (editorial holdings summary), or ''."""
    if not text:
        return ""
    start = text.find(_RATIO_MARKER)
    if start == -1:
        return ""
    body = text[start + len(_RATIO_MARKER):]

    end = len(body)
    dm = _DECISION_START.search(body)
    if dm:
        end = min(end, dm.start())
    for marker in _NEVO_MARKERS:
        if marker == _RATIO_MARKER:
            continue
        pos = body.find(marker)
        if pos != -1:
            end = min(end, pos)
    return body[:end].strip()