# tools/vision_ocr.py - Apple Vision OCR (simple/fast path). No tables or formatting.
from __future__ import annotations

import json
import math
import os
import re
import time

from langchain_core.tools import tool
from pdf2image import convert_from_path
from PIL import Image

from tools._logging import _debug_log, _emit_input_status
from tools.job_stream_schema import (
    close_job_stream,
    error_job_stream,
    make_capabilities,
    metric_def,
    open_job_stream,
    update_job_stream,
)
from tools.live_view_visuals import emit_live_view_preview, image_to_preview_base64
from tools.ocr_backends.vision import VISION_AVAILABLE, VISION_ERROR, run_vision_ocr_on_pil_image
from tools.plugin_stream_controls import get_action_count
from tools.pdf2image_errors import message_for_pdf2image_error
from utils.path_safety import is_dangerous_path, is_plugins_dir_write_blocked

# For ocr_scanned_pdf_to_file: split PDF when larger than this (MB)
_LARGE_FILE_MB = 30


def _is_probably_pdf(path: str) -> bool:
    if (path or "").lower().endswith(".pdf"):
        return True
    try:
        with open(path, "rb") as f:
            return f.read(5) == b"%PDF-"
    except Exception:
        return False


def _parse_split_pdf_result(result: str) -> list[str]:
    """Extract part paths from split_pdf return text (e.g. '1. /path/part1.pdf (123 bytes)\\n2. ...')."""
    paths = []
    for line in result.strip().splitlines():
        line = line.strip()
        if re.match(r"^\d+\.\s+", line):
            rest = line.split(". ", 1)[1].strip()
            if " (" in rest:
                rest = rest.rsplit(" (", 1)[0].strip()
            if rest.endswith(".pdf"):
                paths.append(rest)
    return paths


@tool
def ocr_scanned_pdf(file_path: str) -> str:
    """Fast OCR for scanned PDFs using Apple Vision (macOS). Returns plain text only—no tables or formatting. Use for simple documents; prefer ocr_complex_document for tables, footnotes, or formatted layout."""
    stream_id = f"ocr_vision:{int(time.time() * 1000)}"
    open_job_stream(
        plugin_id="ocr_job",
        plugin_name="OCR Job",
        title="Vision OCR",
        job_type="ocr_scanned_pdf",
        stream_id=stream_id,
        controls=[
            {"id": "cancel", "type": "button", "label": "Cancel OCR", "action_id": "cancel_job", "style": "destructive"}
        ],
        capabilities=make_capabilities(cancel=True, progress=True, metrics=True, phases=True, batch_items=True, artifacts=True),
        metric_schema=[
            metric_def("pages_total", "Pages Total", format="number"),
            metric_def("pages_done", "Pages Done", format="number"),
            metric_def("chars_total", "Characters", format="number"),
        ],
        metadata={"file_path": file_path},
    )
    if not VISION_AVAILABLE:
        msg = VISION_ERROR or "Vision OCR requires macOS."
        error_job_stream(plugin_id="ocr_job", plugin_name="OCR Job", stream_id=stream_id, phase="init", message=msg)
        return msg
    file_path = file_path.strip(' "\n,')
    abs_path = os.path.abspath(os.path.expanduser(file_path))
    if not os.path.exists(abs_path) or not os.path.isfile(abs_path):
        msg = f"File not found: {abs_path}"
        error_job_stream(plugin_id="ocr_job", plugin_name="OCR Job", stream_id=stream_id, phase="init", message=msg)
        return msg
    if not _is_probably_pdf(abs_path):
        msg = "This tool is PDF-only. For PNG/JPG/WebP images, use ocr_image_text(file_path=...)."
        error_job_stream(plugin_id="ocr_job", plugin_name="OCR Job", stream_id=stream_id, phase="init", message=msg)
        return msg
    if is_dangerous_path(abs_path):
        msg = "That path is not allowed for reading."
        error_job_stream(plugin_id="ocr_job", plugin_name="OCR Job", stream_id=stream_id, phase="init", message=msg)
        return msg
    try:
        images = convert_from_path(abs_path, dpi=300)
    except Exception as e:
        msg = message_for_pdf2image_error(e)
        error_job_stream(plugin_id="ocr_job", plugin_name="OCR Job", stream_id=stream_id, phase="convert_pdf", message=msg)
        return msg
    if not images:
        close_job_stream(plugin_id="ocr_job", plugin_name="OCR Job", stream_id=stream_id, status="completed", message="No pages could be converted from PDF.")
        return "No pages could be converted from the PDF."
    full_pages = []
    total = len(images)
    update_job_stream(
        plugin_id="ocr_job",
        plugin_name="OCR Job",
        stream_id=stream_id,
        phase="ocr_pages",
        status="running",
        progress_pct=0,
        message=f"Starting Vision OCR for {total} pages...",
        metrics={"pages_total": float(total), "pages_done": 0.0, "chars_total": 0.0},
    )
    total_chars = 0
    for page_num, pil_img in enumerate(images, 1):
        if get_action_count(stream_id, "cancel_job") > 0:
            close_job_stream(
                plugin_id="ocr_job",
                plugin_name="OCR Job",
                stream_id=stream_id,
                status="stopped",
                message=f"Vision OCR cancelled at page {page_num}/{total}.",
                metrics={"pages_total": float(total), "pages_done": float(page_num - 1), "chars_total": float(total_chars)},
            )
            return "Vision OCR cancelled by user."
        _emit_input_status(f"Vision OCR page {page_num}/{total}…")
        text = run_vision_ocr_on_pil_image(pil_img)
        if not text:
            msg = (
                "Vision OCR bridge is unavailable. Use ocr_complex_document for this file, "
                "or restart Mugi to restore the macOS bridge."
            )
            error_job_stream(plugin_id="ocr_job", plugin_name="OCR Job", stream_id=stream_id, phase="ocr_pages", message=msg)
            return msg
        full_pages.append(text)
        total_chars += len(text or "")
        update_job_stream(
            plugin_id="ocr_job",
            plugin_name="OCR Job",
            stream_id=stream_id,
            phase="ocr_pages",
            status="running",
            progress_pct=float(100 * page_num / total),
            message=f"OCR page {page_num}/{total}",
            metrics={"pages_total": float(total), "pages_done": float(page_num), "chars_total": float(total_chars)},
        )
        try:
            b64 = image_to_preview_base64(pil_img, max_width=320)
            emit_live_view_preview(
                f"OCR page {page_num}/{total}",
                tool_name="ocr_scanned_pdf",
                progress_pct=int(100 * page_num / total),
                preview_image_base64=b64,
                preview_text=text[:300] if text else None,
            )
        except Exception:
            pass
    result = "\n\n".join(full_pages)
    close_job_stream(
        plugin_id="ocr_job",
        plugin_name="OCR Job",
        stream_id=stream_id,
        status="completed",
        message=f"Vision OCR complete ({total} pages).",
        metrics={"pages_total": float(total), "pages_done": float(total), "chars_total": float(total_chars)},
    )
    return result


@tool
def ocr_scanned_pdf_to_file(
    pdf_path: str,
    output_file_path: str,
    append: bool = False,
) -> str:
    """Run Vision OCR on a PDF and write the result to a text file. Plain text only (no tables/formatting). For tables or footnotes use ocr_complex_document_to_file."""
    if not VISION_AVAILABLE:
        return VISION_ERROR or "Vision OCR requires macOS."
    from tools.pdf_tools import split_pdf
    pdf_path = pdf_path.strip(' "\n,')
    output_file_path = output_file_path.strip(' "\n,')
    pdf_abs = os.path.abspath(os.path.expanduser(pdf_path))
    out_abs = os.path.abspath(os.path.expanduser(output_file_path))
    if not os.path.exists(pdf_abs) or not os.path.isfile(pdf_abs):
        return f"PDF not found or not a file: {pdf_abs}"
    if is_dangerous_path(out_abs):
        return f"Output path is in a protected location: {out_abs}"
    if is_plugins_dir_write_blocked(out_abs):
        return "Blocked: cannot write to the plugins directory. Use create_plugin (evolution mode) to add plugins."
    from utils.user_model import enforce_low_level_user_model_write_guard

    _um_err = enforce_low_level_user_model_write_guard(
        out_abs, mode="append" if append else "overwrite"
    )
    if _um_err:
        return _um_err
    file_size_mb = os.path.getsize(pdf_abs) / (1024 * 1024)
    part_paths = []
    if file_size_mb >= _LARGE_FILE_MB:
        num_parts = max(2, math.ceil(file_size_mb / _LARGE_FILE_MB))
        _debug_log(f"[ocr_scanned_pdf_to_file] PDF {file_size_mb:.1f} MB >= {_LARGE_FILE_MB} MB; splitting into {num_parts} parts", "INFO")
        try:
            split_result = split_pdf.invoke({
                "input_json": json.dumps({
                    "file_path": pdf_abs,
                    "split_method": "equal",
                    "num_parts": num_parts,
                }),
            })
            if isinstance(split_result, dict):
                split_result = split_result.get("output", str(split_result))
            part_paths = _parse_split_pdf_result(str(split_result))
        except Exception as e:
            return f"Split failed: {e}"
        if not part_paths:
            return "Split did not return part paths."
    else:
        part_paths = [pdf_abs]
    parent_dir = os.path.dirname(out_abs) or "."
    os.makedirs(parent_dir, exist_ok=True)
    separator = "\n\n---\n\n"
    total_chars = 0
    for i, part_path in enumerate(part_paths):
        _emit_input_status(f"Vision OCR part {i + 1}/{len(part_paths)}…")
        text = ocr_scanned_pdf.invoke({"file_path": part_path})
        if not text or not isinstance(text, str):
            return f"OCR part {i + 1} returned no text."
        if text.startswith("Failed") or text.startswith("File not found") or text.startswith("That path"):
            return f"OCR part {i + 1} failed: {text[:400]}"
        mode = "w" if (i == 0 and not append) else "a"
        with open(out_abs, mode, encoding="utf-8") as f:
            if i > 0 or append:
                f.write(separator)
            f.write(text)
        total_chars += len(text)
    size_bytes = os.path.getsize(out_abs)
    return f"Done. Vision OCR: {len(part_paths)} part(s), {total_chars:,} characters, saved to:\n{out_abs}\n(File size: {size_bytes:,} bytes)."


@tool
def ocr_image_text(file_path: str) -> str:
    """OCR text from an image file (PNG/JPG/JPEG/WebP/TIFF) via the macOS Vision bridge."""
    if not VISION_AVAILABLE:
        return VISION_ERROR or "Vision OCR bridge requires macOS."
    path = (file_path or "").strip(' "\n,')
    if not path:
        return "No file path provided."
    abs_path = os.path.abspath(os.path.expanduser(path))
    if not os.path.exists(abs_path) or not os.path.isfile(abs_path):
        return f"File not found: {abs_path}"
    if is_dangerous_path(abs_path):
        return "That path is not allowed for reading."
    if _is_probably_pdf(abs_path):
        return "This tool is for images only. For PDFs, use ocr_scanned_pdf or ocr_complex_document."
    try:
        with Image.open(abs_path) as img:
            pil_img = img.convert("RGB")
    except Exception as e:
        return f"Failed to open image: {e}"
    text = run_vision_ocr_on_pil_image(pil_img)
    if not text:
        return (
            "Vision OCR bridge is unavailable for image OCR. "
            "Ensure MUGI_CALENDAR_URL and MUGI_AUTH_TOKEN are set, then try again."
        )
    return text
