# /home/Autopartes/pos/services/invoice_ocr.py
"""Supplier-invoice OCR / parsing.

Accepts an image or PDF of a supplier invoice and returns structured lines:

    [{part_number, name, quantity, unit_cost}]

Pipeline (first stage that yields lines wins):
  1. PDF with a text layer  -> ``pdftotext -layout``
  2. Scanned PDF            -> ``pdftoppm`` (first page) + Tesseract
  3. Image                  -> Tesseract
  4. Vision fallback        -> ``ai_chat`` (QWEN / OpenRouter vision model)

Stock is NEVER touched here: the caller returns the parsed lines so the
front-end can confirm before creating a purchase order.

Public API:
    parse_invoice(data, filename=None, content_type=None, use_vision=True) -> dict
"""

import base64
import io
import json
import os
import re
import shutil
import subprocess
import tempfile

_MONEY_RE = re.compile(r"^\$?\s?\d{1,3}(?:[,\d]{0,15})(?:\.\d{1,2})?$")
_INT_RE = re.compile(r"^\d{1,4}$")
_PN_RE = re.compile(r"^[A-Z0-9][A-Z0-9\-\./_]{2,}$")

# Lines that are clearly headers/totals, not invoice items.
_SKIP_TOKENS = {
    "factura", "remision", "total", "subtotal", "iva", "ieps", "rfc", "fecha",
    "pedido", "cliente", "proveedor", "telefono", "direccion", "calle",
    "colonia", "cantidad", "cant", "descripcion", "precio", "importe",
    "unidad", "folio", "orden", "pagina", "gracias", "nota", "clave",
    "codigo", "producto", "articulo", "descuento", "moneda", "condiciones",
    "subtotal:", "total:", "iva:", "rfc:", "fecha:", "importe:", "precio:",
}

# Any of these as a whole word marks a header/footer line (RFC, date, totals).
_HEADER_WORDS = {
    "rfc", "fecha", "folio", "cliente", "proveedor", "factura", "remision",
    "remisión", "pedido", "telefono", "tel", "direccion", "dirección", "calle",
    "colonia", "condiciones", "moneda", "subtotal", "total", "iva", "ieps",
    "clave", "descripcion", "descripción", "cantidad", "unidad", "descuento",
    "importe", "precio", "page",
}

_VISION_PROMPT = (
    "Eres un experto en facturas de refaccionarias mexicanas. Extrae TODAS las "
    "lineas de producto de la imagen. Devuelve SOLO JSON valido, sin texto extra, "
    "con esta forma exacta: "
    '{"lines":[{"part_number":"...","name":"...","quantity":1,"unit_cost":0.0}]}. '
    "part_number = clave/SKU/numero de parte; name = descripcion; quantity = "
    "cantidad (entero); unit_cost = precio unitario (numero, sin simbolos). "
    "Si no hay lineas claras devuelve {\"lines\": []}."
)


# ── Binary helpers ─────────────────────────────────────────────────────────

def _is_pdf(data, filename):
    if data[:5] == b"%PDF-":
        return True
    return bool(filename and filename.lower().endswith(".pdf"))


def _pdftotext(data):
    """Extract the text layer of a PDF. Returns '' on failure."""
    exe = shutil.which("pdftotext")
    if not exe:
        return ""
    try:
        proc = subprocess.run(
            [exe, "-layout", "-", "-"],
            input=data, stdout=subprocess.PIPE, stderr=subprocess.PIPE, timeout=60,
        )
        if proc.returncode == 0:
            return proc.stdout.decode("utf-8", "replace")
    except Exception:
        pass
    return ""


def _pdf_first_page_png(data, dpi=200):
    """Render the first PDF page to PNG bytes (for OCR). None on failure."""
    exe = shutil.which("pdftoppm")
    if not exe:
        return None
    try:
        with tempfile.TemporaryDirectory() as td:
            pdf_path = os.path.join(td, "in.pdf")
            with open(pdf_path, "wb") as f:
                f.write(data)
            prefix = os.path.join(td, "page")
            subprocess.run(
                [exe, "-png", "-r", str(dpi), "-f", "1", "-l", "1", pdf_path, prefix],
                stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, timeout=120,
            )
            for fn in sorted(os.listdir(td)):
                if fn.startswith("page") and fn.endswith(".png"):
                    with open(os.path.join(td, fn), "rb") as f:
                        return f.read()
    except Exception:
        pass
    return None


def _ocr_langs():
    try:
        import pytesseract
        langs = pytesseract.get_languages(config="")
        for wanted in ("spa", "eng"):
            if wanted in langs:
                return wanted
    except Exception:
        pass
    return "eng"


def _ocr_image(img_bytes):
    """Run Tesseract over image bytes. Returns '' on failure."""
    try:
        from PIL import Image
        import pytesseract
    except Exception:
        return ""
    try:
        img = Image.open(io.BytesIO(img_bytes))
        if img.mode not in ("RGB", "L"):
            img = img.convert("RGB")
        return pytesseract.image_to_string(img, lang=_ocr_langs())
    except Exception:
        return ""


def _ocr_rows(img_bytes, y_tolerance=12):
    """OCR an image and rebuild table rows from word positions.

    Tesseract's default layout analysis splits a monospaced invoice table into
    one block per column, so ``image_to_string`` returns each column stacked
    instead of each row.  We regroup words by their vertical center instead.

    Returns a list of row strings (falls back to plain lines on failure).
    """
    try:
        from PIL import Image
        import pytesseract
    except Exception:
        return []
    try:
        img = Image.open(io.BytesIO(img_bytes))
        if img.mode not in ("RGB", "L"):
            img = img.convert("RGB")
        data = pytesseract.image_to_data(
            img, lang=_ocr_langs(), output_type=pytesseract.Output.DICT
        )
    except Exception:
        text = _ocr_image(img_bytes)
        return text.splitlines() if text else []

    groups = {}  # representative y -> list of (left, text)
    for i, txt in enumerate(data.get("text", [])):
        txt = (txt or "").strip()
        if not txt:
            continue
        try:
            top = int(data["top"][i])
            h = int(data["height"][i] or 0)
            left = int(data["left"][i])
        except (KeyError, ValueError, TypeError):
            continue
        cy = top + h // 2
        key = None
        for k in groups:
            if abs(k - cy) <= max(y_tolerance, h // 2):
                key = k
                break
        groups.setdefault(cy if key is None else key, []).append((left, txt))

    rows = []
    for cy in sorted(groups):
        rows.append(" ".join(t for _, t in sorted(groups[cy])))
    return rows


def _ocr_rows_text(img_bytes):
    """Row-reconstructed OCR text (falls back to raw Tesseract output)."""
    rows = _ocr_rows(img_bytes)
    if rows:
        return "\n".join(rows)
    return _ocr_image(img_bytes)


# ── Text line parser ───────────────────────────────────────────────────────

def _is_money(tok):
    if tok.startswith("$"):
        return True
    if not _MONEY_RE.match(tok):
        return False
    return ("." in tok) or ("," in tok)


def _to_money(tok):
    cleaned = re.sub(r"[^\d.]", "", tok.replace(",", ""))
    try:
        return float(cleaned)
    except ValueError:
        return None


def _is_part_number(tok):
    if _is_money(tok):
        return False
    if not _PN_RE.match(tok):
        return False
    has_alpha = any(c.isalpha() for c in tok)
    digits = sum(c.isdigit() for c in tok)
    if has_alpha:
        return digits >= 1 and len(tok) >= 4
    # pure numeric part numbers are accepted only if long enough to not be a qty
    return len(tok) >= 4


def _skip_line(line):
    low = line.lower()
    if low.strip(".:-") in _SKIP_TOKENS:
        return True
    words = set(re.findall(r"[a-záéíóúñ]+", low))
    if words & _HEADER_WORDS:
        return True
    if words and all(t in _SKIP_TOKENS or t in ("de", "del", "la", "el") for t in words):
        return True
    return False


def _parse_line(line):
    tokens = line.split()
    if len(tokens) < 3:
        return None

    money_idx = [i for i, t in enumerate(tokens) if _is_money(t)]
    int_idx = [i for i, t in enumerate(tokens) if _INT_RE.match(t)]
    pn_idx = [i for i, t in enumerate(tokens) if _is_part_number(t)]
    if not pn_idx or (not money_idx and not int_idx):
        return None

    qty, qty_i = None, None
    for i in int_idx:
        if i in money_idx:
            continue
        v = int(tokens[i])
        if 1 <= v <= 9999:
            qty, qty_i = v, i
            break

    unit, unit_i = None, None
    if money_idx:
        unit_i = money_idx[0]
        unit = _to_money(tokens[unit_i])

    pn_i = pn_idx[0]
    used = {i for i in (pn_i, qty_i) if i is not None} | set(money_idx)
    name = " ".join(t for i, t in enumerate(tokens) if i not in used).strip(" -|")
    if not name:
        name = None
    if qty is None and unit is None:
        return None
    return {
        "part_number": tokens[pn_i],
        "name": name,
        "quantity": qty,
        "unit_cost": unit,
    }


def parse_lines(text):
    """Parse raw invoice text into structured lines."""
    out = []
    seen = set()
    for raw in (text or "").splitlines():
        line = raw.strip()
        if len(line) < 6 or _skip_line(line):
            continue
        row = _parse_line(line)
        if not row:
            continue
        key = (row["part_number"].upper(), row["quantity"], row["unit_cost"])
        if key in seen:
            continue
        seen.add(key)
        out.append(row)
    return out


# ── Vision fallback ────────────────────────────────────────────────────────

def _extract_json(text):
    if not text:
        return None
    s = text.strip()
    if s.startswith("```"):
        s = re.sub(r"^```[a-zA-Z]*\n?", "", s)
        s = re.sub(r"\n?```$", "", s)
    try:
        return json.loads(s)
    except Exception:
        pass
    m = re.search(r"\{.*\}", s, re.DOTALL)
    if m:
        try:
            return json.loads(m.group(0))
        except Exception:
            return None
    return None


def _coerce_lines(obj):
    if not obj:
        return []
    raw = obj.get("lines") if isinstance(obj, dict) else None
    if raw is None and isinstance(obj, list):
        raw = obj
    if not isinstance(raw, list):
        return []
    lines = []
    for item in raw:
        if not isinstance(item, dict):
            continue
        pn = item.get("part_number") or item.get("clave") or item.get("sku")
        name = item.get("name") or item.get("descripcion") or item.get("description")
        qty = item.get("quantity", item.get("cantidad"))
        cost = item.get("unit_cost", item.get("precio", item.get("unit_price")))
        try:
            qty = int(qty) if qty is not None else None
        except (ValueError, TypeError):
            qty = None
        try:
            cost = float(cost) if cost is not None else None
        except (ValueError, TypeError):
            cost = None
        if pn or name:
            lines.append({
                "part_number": str(pn).strip() if pn else None,
                "name": str(name).strip() if name else None,
                "quantity": qty, "unit_cost": cost,
            })
    return lines


def _vision_extract(image_bytes, hint_text=None):
    """Extract lines with the vision model. Returns (lines, error)."""
    try:
        from services import ai_chat
    except Exception as e:
        return [], f"ai_chat import: {e}"

    b64 = base64.b64encode(image_bytes).decode("ascii")
    prompt = _VISION_PROMPT
    if hint_text:
        prompt += "\n\nTexto OCR parcial (puede estar incompleto):\n" + hint_text[:1500]

    # 1) reuse chat_with_image (per task requirement)
    try:
        res = ai_chat.chat_with_image(prompt, b64)
        lines = _coerce_lines(res)
        if lines:
            return lines, None
        msg = res.get("message") if isinstance(res, dict) else None
        lines = _coerce_lines(_extract_json(msg) if msg else None)
        if lines:
            return lines, None
    except Exception:
        pass

    # 2) direct vision call with an invoice-specific system prompt
    try:
        data_url = "data:image/jpeg;base64," + b64
        messages = [
            {"role": "system", "content": "Responde SOLO JSON valido."},
            {"role": "user", "content": [
                {"type": "image_url", "image_url": {"url": data_url}},
                {"type": "text", "text": prompt},
            ]},
        ]
        backends = []
        if ai_chat.QWEN_ENABLED:
            backends.append((ai_chat.QWEN_CHAT_URL, ai_chat.QWEN_API_KEY, ai_chat.QWEN_MODEL))
        if ai_chat.OPENROUTER_API_KEY:
            backends.append((ai_chat.OPENROUTER_URL, ai_chat.OPENROUTER_API_KEY, ai_chat.VISION_MODEL))
        for url, key, model in backends:
            result = ai_chat._post_chat_completion(
                url, key, model, messages, max_tokens=1200, temperature=0.1, timeout=60
            )
            if not result:
                continue
            lines = _coerce_lines(_extract_json(result["content"]))
            if lines:
                return lines, None
        return [], "vision sin lineas"
    except Exception as e:
        return [], f"vision: {e}"


# ── Public API ─────────────────────────────────────────────────────────────

def parse_invoice(data, filename=None, content_type=None, use_vision=True):
    """Parse an invoice image/PDF into structured lines.

    Returns:
        {
          "lines": [{part_number, name, quantity, unit_cost}],
          "count": int,
          "method": str,           # pdftotext | pdftoppm+tesseract | tesseract | ...+vision
          "warnings": [str],
          "text_preview": str,
        }
    """
    warnings = []
    text = ""
    method = "unknown"
    vision_image = None  # bytes to hand to the vision model when OCR fails

    if not data:
        return {"lines": [], "count": 0, "method": "empty",
                "warnings": ["archivo vacío"], "text_preview": ""}

    if _is_pdf(data, filename):
        text = _pdftotext(data)
        if len(text.strip()) >= 30:
            method = "pdftotext"
        else:
            png = _pdf_first_page_png(data)
            if png:
                vision_image = png
                text = _ocr_rows_text(png)
                method = "pdftoppm+tesseract" if text.strip() else "pdftoppm"
                if not text.strip():
                    warnings.append("PDF sin capa de texto y OCR vacío")
            else:
                method = "pdf"
                warnings.append("No se pudo renderizar el PDF")
        if not vision_image:
            # for vision fallback on a text PDF we still need a raster copy
            vision_image = _pdf_first_page_png(data)
    else:
        vision_image = data
        text = _ocr_rows_text(data)
        method = "tesseract"
        if not text.strip():
            warnings.append("OCR sin texto reconocible")

    lines = parse_lines(text)

    if not lines and use_vision and vision_image:
        vlines, verr = _vision_extract(vision_image, hint_text=text)
        if vlines:
            lines = vlines
            method += "+vision"
        elif verr:
            warnings.append(f"fallback visión sin resultados: {verr}")

    return {
        "lines": lines,
        "count": len(lines),
        "method": method,
        "warnings": warnings,
        "text_preview": text[:800],
    }