#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
convert_ocr.py

OCR専用の独立変換ツール。
- 出力は常に Markdown (.md)
- OCRエンジンは tesseract / genai
- 入力は画像, PDF, および (LibreOffice/soffice があれば) Office 文書
- --list で利用可能な形式と依存関係を表示

主な使い方:
    python convert_ocr.py input.png
    python convert_ocr.py input.pdf --engine tesseract --lang jpn+eng
    python convert_ocr.py input.png --engine genai --ini ai_ocr2md.ini
    python convert_ocr.py slides.pptx --engine tesseract --keep-images
    python convert_ocr.py --list
"""

from __future__ import annotations

import argparse
import os
import re
import shutil
import subprocess
import sys
import tempfile
from dataclasses import dataclass
from pathlib import Path
from typing import Dict, List, Optional, Sequence, Tuple


# ---------- optional imports ----------
try:
    from PIL import Image, ImageGrab
except Exception:  # pragma: no cover
    Image = None
    ImageGrab = None

try:
    import fitz  # PyMuPDF
except Exception:  # pragma: no cover
    fitz = None

try:
    import pytesseract
except Exception:  # pragma: no cover
    pytesseract = None


# ---------- constants ----------
IMAGE_EXTS = {".png", ".jpg", ".jpeg", ".bmp", ".tif", ".tiff", ".webp"}
PDF_EXTS = {".pdf"}
OFFICE_EXTS = {
    ".pptx", ".ppt", ".docx", ".doc", ".xlsx", ".xls", ".odp", ".odt", ".ods"
}
DEFAULT_GENAI_PROMPT = (
    "画像内のテキストを正確に抽出し、Markdown形式で出力してください。"
    " 見出し・箇条書き・表をできるだけ保ち、数式は LaTeX で表現してください。"
)


# ---------- data structures ----------
@dataclass
class ImageItem:
    path: Path
    label: str


# ---------- utility ----------
def eprint(*args: object, **kwargs: object) -> None:
    print(*args, file=sys.stderr, **kwargs)


def which(cmd: str) -> Optional[str]:
    return shutil.which(cmd)


def command_exists_any(names: Sequence[str]) -> Optional[str]:
    for name in names:
        found = which(name)
        if found:
            return found
    return None


def slugify_filename(name: str) -> str:
    return re.sub(r"[^A-Za-z0-9._-]+", "_", name)


def default_output_path(input_value: str, output: Optional[str]) -> Path:
    if output:
        return Path(output)
    if input_value.lower() == "clip":
        return Path("clipboard_ocr.md")
    return Path(input_value).with_suffix(".md")


# ---------- ini reader ----------
def read_simple_ini(path: Path) -> Dict[str, str]:
    """
    key=value 形式の簡易 INI リーダー。
    # / ; コメント、空行、三連引用符複数行値に対応。
    """
    result: Dict[str, str] = {}
    text = path.read_text(encoding="utf-8")
    lines = text.splitlines()
    i = 0
    while i < len(lines):
        line = lines[i].strip()
        i += 1
        if not line or line.startswith("#") or line.startswith(";"):
            continue
        if "=" not in line:
            continue
        key, value = line.split("=", 1)
        key = key.strip()
        value = value.strip()

        if value.startswith("'''") or value.startswith('"""'):
            quote = value[:3]
            if value.endswith(quote) and len(value) >= 6:
                result[key] = value[3:-3]
                continue
            collected = [value[3:]]
            while i < len(lines):
                cur = lines[i]
                i += 1
                if cur.endswith(quote):
                    collected.append(cur[:-3])
                    break
                collected.append(cur)
            result[key] = "\n".join(collected)
        else:
            result[key] = value

    # simple environment expansion: $VAR
    env_pattern = re.compile(r"\$([A-Za-z_][A-Za-z0-9_]*)")
    for key, value in list(result.items()):
        result[key] = env_pattern.sub(lambda m: os.getenv(m.group(1), m.group(0)), value)
    return result


# ---------- availability / list ----------
def availability_report() -> str:
    soffice = command_exists_any(["soffice", "libreoffice"])
    lines = []
    lines.append("convert_ocr.py : 利用可能な入力 / エンジン / 依存関係")
    lines.append("")
    lines.append("[OCR engine]")
    lines.append(f"- tesseract : {'OK' if pytesseract is not None else 'NG'}"
                 f"{' (pytesseract import可)' if pytesseract is not None else ' (pytesseract が必要)'}")
    lines.append(f"- genai     : {'OK' if probe_genai_backend()[0] else 'NG'}"
                 f" ({probe_genai_backend()[1]})")
    lines.append("")
    lines.append("[renderer]")
    lines.append(f"- Pillow    : {'OK' if Image is not None else 'NG'}")
    lines.append(f"- PyMuPDF   : {'OK' if fitz is not None else 'NG'}")
    lines.append(f"- soffice   : {'OK' if soffice else 'NG'}"
                 f"{' (' + soffice + ')' if soffice else ' (LibreOffice があれば Office 文書入力可)'}")
    lines.append("")
    lines.append("[input type]")
    lines.append(f"- image     : {', '.join(sorted(IMAGE_EXTS))}")
    lines.append(f"- pdf       : {', '.join(sorted(PDF_EXTS))}")
    lines.append(f"- office    : {', '.join(sorted(OFFICE_EXTS))}  ※ soffice が必要")
    lines.append("- clipboard : clip  ※ Pillow.ImageGrab が使える環境のみ")
    lines.append("")
    lines.append("[output]")
    lines.append("- 常に .md")
    lines.append("")
    lines.append("[examples]")
    lines.append("- python convert_ocr.py input.png")
    lines.append("- python convert_ocr.py input.pdf --engine tesseract --lang jpn+eng")
    lines.append("- python convert_ocr.py input.png --engine genai --ini ai_ocr2md.ini")
    lines.append("- python convert_ocr.py slides.pptx --engine tesseract --keep-images")
    return "\n".join(lines)


# ---------- genai backend ----------
def probe_genai_backend() -> Tuple[bool, str]:
    try:
        import tkai_lib  # type: ignore
        return True, "tkai_lib"
    except Exception:
        pass
    try:
        import google.generativeai  # type: ignore
        return True, "google.generativeai"
    except Exception:
        pass
    return False, "tkai_lib または google-generativeai が必要"


def load_genai_prompt(ini_path: Optional[str]) -> str:
    if not ini_path:
        return DEFAULT_GENAI_PROMPT
    path = Path(ini_path)
    if not path.exists():
        eprint(f"[warn] ini file not found: {path}. 既定プロンプトを使用します。")
        return DEFAULT_GENAI_PROMPT
    data = read_simple_ini(path)
    return (
        data.get("PROMPT")
        or data.get("PROMPT_TEMPLATE_JA")
        or data.get("PROMPT_TEMPLATE_EN")
        or DEFAULT_GENAI_PROMPT
    )


def init_genai_environment(ai_env: Optional[str]) -> None:
    if not ai_env:
        return
    env_path = Path(ai_env)
    if not env_path.exists():
        return
    data = read_simple_ini(env_path)
    for key, value in data.items():
        if key and value and key not in os.environ:
            os.environ[key] = value


def call_genai_ocr(prompt: str, image_path: Path, api: str = "gemini", model: Optional[str] = None,
                   ai_env: Optional[str] = None) -> str:
    if Image is None:
        raise RuntimeError("Pillow is required for genai OCR.")

    init_genai_environment(ai_env)
    img = Image.open(image_path)

    # Backend 1: user's existing wrapper
    try:
        import tkai_lib  # type: ignore
        if hasattr(tkai_lib, "read_ai_config") and ai_env:
            tkai_lib.read_ai_config(ai_env)
        if api != "gemini":
            raise RuntimeError("この実装の genai は現在 gemini を想定しています。")
        model_name = model or os.getenv("gemini_model", "gemini-1.5-pro")
        api_key = os.getenv("GOOGLE_API_KEY")
        if not api_key:
            raise RuntimeError("GOOGLE_API_KEY is not set.")
        tkai_lib.genai.configure(api_key=api_key)
        m = tkai_lib.genai.GenerativeModel(model_name)
        response = m.generate_content([prompt, img])
        text = getattr(response, "text", None)
        if not text:
            raise RuntimeError("genai response.text is empty.")
        return text.strip()
    except ImportError:
        pass

    # Backend 2: direct google.generativeai
    try:
        import google.generativeai as genai  # type: ignore
        if api != "gemini":
            raise RuntimeError("この実装の genai は現在 gemini を想定しています。")
        model_name = model or os.getenv("gemini_model", "gemini-1.5-pro")
        api_key = os.getenv("GOOGLE_API_KEY")
        if not api_key:
            raise RuntimeError("GOOGLE_API_KEY is not set.")
        genai.configure(api_key=api_key)
        m = genai.GenerativeModel(model_name)
        response = m.generate_content([prompt, img])
        text = getattr(response, "text", None)
        if not text:
            raise RuntimeError("genai response.text is empty.")
        return text.strip()
    except ImportError as exc:
        raise RuntimeError("genai backend not available. Install tkai_lib or google-generativeai.") from exc


# ---------- input normalization ----------
def ensure_clipboard_image(work_dir: Path) -> Path:
    if ImageGrab is None:
        raise RuntimeError("Pillow.ImageGrab is not available in this environment.")
    img = ImageGrab.grabclipboard()
    if not isinstance(img, Image.Image):
        raise RuntimeError("Clipboard does not contain an image.")
    out = work_dir / "clipboard_input.png"
    img.save(out, "PNG")
    return out


def render_pdf_to_images(pdf_path: Path, out_dir: Path, dpi: int) -> List[ImageItem]:
    if fitz is None:
        raise RuntimeError("PyMuPDF (fitz) is required for PDF input.")
    out_dir.mkdir(parents=True, exist_ok=True)
    items: List[ImageItem] = []
    doc = fitz.open(pdf_path)
    try:
        matrix = fitz.Matrix(dpi / 72.0, dpi / 72.0)
        for page_idx in range(len(doc)):
            page = doc.load_page(page_idx)
            pix = page.get_pixmap(matrix=matrix, alpha=False)
            out_path = out_dir / f"page_{page_idx + 1:04d}.png"
            pix.save(str(out_path))
            items.append(ImageItem(path=out_path, label=f"Page {page_idx + 1}"))
    finally:
        doc.close()
    return items


def convert_office_to_pdf(input_path: Path, out_dir: Path) -> Path:
    soffice = command_exists_any(["soffice", "libreoffice"])
    if not soffice:
        raise RuntimeError(
            "Office文書入力には LibreOffice/soffice が必要です。"
            " 先に PDF に変換するか、LibreOffice をインストールしてください。"
        )
    out_dir.mkdir(parents=True, exist_ok=True)
    cmd = [
        soffice,
        "--headless",
        "--convert-to", "pdf",
        "--outdir", str(out_dir),
        str(input_path),
    ]
    proc = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
    if proc.returncode != 0:
        raise RuntimeError(
            "soffice conversion failed.\n"
            f"stdout:\n{proc.stdout}\n"
            f"stderr:\n{proc.stderr}"
        )
    pdf_path = out_dir / (input_path.stem + ".pdf")
    if not pdf_path.exists():
        candidates = sorted(out_dir.glob("*.pdf"))
        if not candidates:
            raise RuntimeError("soffice did not produce a PDF file.")
        pdf_path = candidates[0]
    return pdf_path


def collect_images(input_value: str, work_dir: Path, dpi: int) -> List[ImageItem]:
    if input_value.lower() == "clip":
        clip_path = ensure_clipboard_image(work_dir)
        return [ImageItem(path=clip_path, label="Clipboard")]

    input_path = Path(input_value)
    if not input_path.exists():
        raise FileNotFoundError(f"input file not found: {input_path}")

    suffix = input_path.suffix.lower()
    if suffix in IMAGE_EXTS:
        return [ImageItem(path=input_path, label=input_path.name)]

    if suffix in PDF_EXTS:
        return render_pdf_to_images(input_path, work_dir / f"{slugify_filename(input_path.stem)}_images", dpi)

    if suffix in OFFICE_EXTS:
        pdf_dir = work_dir / f"{slugify_filename(input_path.stem)}_pdf"
        pdf_path = convert_office_to_pdf(input_path, pdf_dir)
        return render_pdf_to_images(pdf_path, work_dir / f"{slugify_filename(input_path.stem)}_images", dpi)

    raise RuntimeError(f"unsupported input type: {suffix}")


# ---------- OCR ----------
def tesseract_ocr_image(image_path: Path, lang: str = "jpn+eng") -> str:
    if pytesseract is None:
        raise RuntimeError("pytesseract is required for tesseract OCR.")
    if Image is None:
        raise RuntimeError("Pillow is required for tesseract OCR.")
    img = Image.open(image_path)
    return pytesseract.image_to_string(img, lang=lang).strip()


def build_markdown(items: List[ImageItem], texts: List[str], engine: str, source_name: str) -> str:
    lines: List[str] = []
    lines.append(f"# OCR Result: {source_name}")
    lines.append("")
    lines.append(f"- Engine: `{engine}`")
    lines.append(f"- Segments: {len(items)}")
    lines.append("")

    for idx, (item, text) in enumerate(zip(items, texts), start=1):
        lines.append("---")
        lines.append("")
        lines.append(f"## {item.label}")
        lines.append("")
        if text.strip():
            lines.append(text.rstrip())
        else:
            lines.append("*(empty result)*")
        lines.append("")

    return "\n".join(lines).rstrip() + "\n"


# ---------- main ----------
def parse_args(argv: Optional[Sequence[str]] = None) -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="OCR専用変換ツール。出力は常に Markdown (.md)")
    parser.add_argument("input", nargs="?", default="", help="入力ファイル。画像/PDF/Office文書、または 'clip'")
    parser.add_argument("--output", "-o", default=None, help="出力Markdownファイル名")
    parser.add_argument("--engine", "-e", choices=["tesseract", "genai"], default="tesseract",
                        help="OCRエンジン")
    parser.add_argument("--ini", "-i", default=None, help="genAI用プロンプト設定ファイル")
    parser.add_argument("--api", "-a", default="gemini", help="genAI API名 (既定: gemini)")
    parser.add_argument("--model", "-m", default=None, help="genAIモデル名")
    parser.add_argument("--lang", default="jpn+eng", help="tesseract OCR言語")
    parser.add_argument("--dpi", type=int, default=300, help="PDF/Officeを画像化する際のDPI")
    parser.add_argument("--keep-images", action="store_true",
                        help="中間画像・中間PDFを削除せず残す")
    parser.add_argument("--workdir", default="", help="作業ディレクトリ。未指定時は一時ディレクトリ")
    parser.add_argument("--ai-env", default="ai.env", help="AI用環境変数ファイル (既定: ai.env)")
    parser.add_argument("--list", action="store_true", help="対応形式・依存関係・実行例を表示して終了")
    return parser.parse_args(argv)


def main(argv: Optional[Sequence[str]] = None) -> int:
    args = parse_args(argv)

    if args.list:
        print(availability_report())
        return 0

    if not args.input:
        eprint("error: input is required unless --list is specified")
        return 2

    output_path = default_output_path(args.input, args.output)
    workdir_obj: Optional[tempfile.TemporaryDirectory[str]] = None
    if args.workdir:
        work_dir = Path(args.workdir)
        work_dir.mkdir(parents=True, exist_ok=True)
    else:
        workdir_obj = tempfile.TemporaryDirectory(prefix="convert_ocr_")
        work_dir = Path(workdir_obj.name)

    try:
        items = collect_images(args.input, work_dir, args.dpi)
        texts: List[str] = []

        if args.engine == "tesseract":
            for item in items:
                print(f"[tesseract] OCR: {item.label} -> {item.path}")
                texts.append(tesseract_ocr_image(item.path, lang=args.lang))
        else:
            prompt = load_genai_prompt(args.ini)
            for item in items:
                print(f"[genai] OCR: {item.label} -> {item.path}")
                texts.append(call_genai_ocr(prompt, item.path, api=args.api, model=args.model, ai_env=args.ai_env))

        md_text = build_markdown(items, texts, args.engine, Path(args.input).name if args.input != "clip" else "clipboard")
        output_path.parent.mkdir(parents=True, exist_ok=True)
        output_path.write_text(md_text, encoding="utf-8")
        print(f"OK: wrote {output_path}")

        if args.keep_images:
            print(f"work files kept in: {work_dir}")
            workdir_obj = None  # prevent cleanup of temp dir object below
        return 0

    except Exception as exc:
        eprint(f"ERROR: {exc}")
        return 1

    finally:
        if workdir_obj is not None:
            workdir_obj.cleanup()


if __name__ == "__main__":
    raise SystemExit(main())