# app.py — DeepSeek-OCR (HF Space, Gradio only, recursive finder)
import io, os, sys, base64, importlib.util, tempfile, traceback
from typing import Optional, List
from PIL import Image
import numpy as np
import gradio as gr

ROOT = os.path.dirname(__file__)

TARGET_FILENAMES = [
    "deepseek_ocr.py",          # 함수형/클래스형
    "run_dpsk_ocr_image.py",    # CLI 스타일
    "run_dpsk_ocr.py",          # HF 전용 스크립트
]

def find_candidates(root: str) -> List[str]:
    """프로젝트 전체를 재귀적으로 뒤져 타겟 파일 경로들을 모두 수집"""
    hits = []
    for dirpath, dirnames, filenames in os.walk(root):
        for fn in filenames:
            if fn in TARGET_FILENAMES:
                hits.append(os.path.join(dirpath, fn))
    return sorted(hits)

def load_module_from_path(path: str):
    name = os.path.splitext(os.path.basename(path))[0]
    spec = importlib.util.spec_from_file_location(name, path)
    if not spec or not spec.loader:
        raise ImportError(f"Cannot load module from {path}")
    mod = importlib.util.module_from_spec(spec)
    sys.modules[name] = mod
    spec.loader.exec_module(mod)
    return mod

class OCRAdapter:
    def __init__(self):
        self.mode = "demo"
        self.path = None
        self.debug = []
        self.entry = lambda img, lang="auto": "[DEMO] DeepSeek 코드 연결 전입니다."

        hits = find_candidates(ROOT)
        self.debug.append(f"ROOT={ROOT}")
        self.debug.append("FOUND=" + ("; ".join(hits) if hits else "(none)"))

        for path in hits:
            try:
                mod = load_module_from_path(path)
                # 1) 함수형: ocr_image(image, lang="auto")
                if hasattr(mod, "ocr_image"):
                    self.entry = lambda img, lang="auto", _m=mod: _m.ocr_image(img, lang=lang)
                    self.mode, self.path = "func_ocr_image", path
                    self.debug.append(f"USE {path} :: ocr_image")
                    return
                # 2) 클래스형: DeepSeekOCR().recognize(image, lang)
                if hasattr(mod, "DeepSeekOCR"):
                    inst = mod.DeepSeekOCR()
                    if hasattr(inst, "recognize"):
                        self.entry = lambda img, lang="auto", _i=inst: _i.recognize(img, lang=lang)
                        self.mode, self.path = "class_recognize", path
                        self.debug.append(f"USE {path} :: DeepSeekOCR.recognize")
                        return
                # 3) 스크립트형: run / infer / main (파일경로 요구 가능)
                for cand in ("run", "infer", "main", "predict"):
                    if hasattr(mod, cand):
                        fn = getattr(mod, cand)
                        def _call(img, lang="auto", _fn=fn):
                            with tempfile.NamedTemporaryFile(suffix=".png", delete=True) as tmp:
                                img.save(tmp.name)
                                try:
                                    return str(_fn(tmp.name))
                                except TypeError:
                                    return str(_fn(tmp.name, lang=lang))
                        self.entry = _call
                        self.mode, self.path = f"script_{cand}", path
                        self.debug.append(f"USE {path} :: {cand}(path)")
                        return
                self.debug.append(f"NO ENTRY in {path}")
            except Exception as e:
                self.debug.append(f"LOAD FAIL {path} :: {e}")

    def recognize(self, image: Image.Image, lang="auto") -> str:
        return self.entry(image.convert("RGB"), lang)

ADAPTER = OCRAdapter()

def _to_pil(x) -> Image.Image:
    if isinstance(x, Image.Image): return x.convert("RGB")
    if isinstance(x, (bytes, bytearray)): return Image.open(io.BytesIO(x)).convert("RGB")
    if isinstance(x, np.ndarray): return Image.fromarray(x).convert("RGB")
    raise TypeError("Unsupported image type")

def _b64_to_image(image_b64: str) -> Image.Image:
    return _to_pil(base64.b64decode(image_b64))

def gradio_predict(image, lang):
    if image is None: return "No image provided."
    try:
        return ADAPTER.recognize(_to_pil(image), lang)
    except Exception as e:
        return f"[ERROR] {e}\n" + traceback.format_exc()

with gr.Blocks(title="DeepSeek-OCR (HF Space, Gradio)") as demo:
    gr.Markdown(
        "### DeepSeek-OCR (HF Space, Gradio)\n"
        f"**현재 모드:** `{ADAPTER.mode}`  \n"
        f"**경로:** `{ADAPTER.path}`  \n"
        f"**찾은 후보:**  \n```\n" + "\n".join(ADAPTER.debug) + "\n```"
    )
    with gr.Row():
        img = gr.Image(type="pil", label="Input Image")
        out = gr.Textbox(label="OCR Result", lines=10)
    lang = gr.Radio(["auto","en","ko","ja","zh"], value="auto", label="Language")
    btn = gr.Button("Run OCR")
    btn.click(gradio_predict, inputs=[img, lang], outputs=[out])

demo.launch()