import json
import os
import tempfile
from typing import Any, Dict, List, Tuple

import gradio as gr
from PIL import Image
import fitz  # PyMuPDF
from paddleocr import PaddleOCR


# =========================
# Config
# =========================

APP_TITLE = os.environ.get("APP_TITLE", "DJ OCR Lab")
APP_SUBTITLE = os.environ.get(
    "APP_SUBTITLE",
    "Upload an image or PDF and run OCR locally. No API. No branding. No cloud goblin."
)

LANG = os.environ.get("OCR_LANG", "en")
PDF_DPI = int(os.environ.get("PDF_DPI", "200"))

# Initialize OCR once when the app starts.
# use_angle_cls helps rotated/skewed text.
ocr = PaddleOCR(
    use_angle_cls=True,
    lang=LANG,
)


# =========================
# File Helpers
# =========================

def is_pdf(path: str) -> bool:
    return path.lower().endswith(".pdf")


def is_image(path: str) -> bool:
    return path.lower().endswith((".png", ".jpg", ".jpeg", ".webp", ".bmp", ".tif", ".tiff"))


def render_pdf_to_images(pdf_path: str, dpi: int = 200) -> List[str]:
    """
    Converts each PDF page to a temporary PNG image.
    """
    image_paths = []

    doc = fitz.open(pdf_path)
    zoom = dpi / 72
    matrix = fitz.Matrix(zoom, zoom)

    for page_index in range(len(doc)):
        page = doc.load_page(page_index)
        pix = page.get_pixmap(matrix=matrix, alpha=False)

        out_path = os.path.join(
            tempfile.gettempdir(),
            f"ocr_page_{os.path.basename(pdf_path)}_{page_index + 1}.png"
        )

        pix.save(out_path)
        image_paths.append(out_path)

    doc.close()
    return image_paths


def normalize_ocr_result(raw_result: Any) -> List[Dict[str, Any]]:
    """
    PaddleOCR return shapes can vary a bit by version.
    This tries to normalize common outputs into:
    [
      {
        "text": "...",
        "confidence": 0.99,
        "box": [...]
      }
    ]
    """
    items = []

    if not raw_result:
        return items

    # Common shape:
    # [
    #   [
    #     [[[x,y],...], ("text", confidence)],
    #     ...
    #   ]
    # ]
    if isinstance(raw_result, list):
        first_layer = raw_result

        # If wrapped as one page, unwrap it.
        if len(first_layer) == 1 and isinstance(first_layer[0], list):
            first_layer = first_layer[0]

        for entry in first_layer:
            try:
                box = entry[0]
                text = entry[1][0]
                confidence = float(entry[1][1])

                items.append({
                    "text": text,
                    "confidence": confidence,
                    "box": box,
                })
            except Exception:
                # Keep weird entries instead of losing evidence.
                items.append({
                    "text": str(entry),
                    "confidence": None,
                    "box": None,
                })

    return items


def ocr_image(image_path: str) -> Tuple[str, List[Dict[str, Any]]]:
    raw = ocr.ocr(image_path, cls=True)
    items = normalize_ocr_result(raw)

    text = "\n".join(item["text"] for item in items if item.get("text"))
    return text.strip(), items


def run_local_ocr(file_path: str) -> Tuple[str, str]:
    if not file_path:
        raise gr.Error("Please upload an image or PDF first.")

    if not os.path.exists(file_path):
        raise gr.Error("Uploaded file was not found.")

    page_outputs = []
    structured_outputs = []

    if is_pdf(file_path):
        image_paths = render_pdf_to_images(file_path, dpi=PDF_DPI)
    elif is_image(file_path):
        image_paths = [file_path]
    else:
        raise gr.Error("Unsupported file type. Please upload an image or PDF.")

    for index, image_path in enumerate(image_paths, start=1):
        text, items = ocr_image(image_path)

        page_outputs.append(f"## Page {index}\n\n{text or '[No text recognized]'}")

        structured_outputs.append({
            "page": index,
            "image_path": image_path,
            "items": items,
        })

    markdown_text = "\n\n---\n\n".join(page_outputs)
    json_output = json.dumps(structured_outputs, indent=2, ensure_ascii=False)

    return markdown_text, json_output


# =========================
# UI
# =========================

custom_css = """
body,
.gradio-container {
    font-family: Inter, system-ui, -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif;
}

.app-header {
    text-align: center;
    max-width: 900px;
    margin: 0 auto 16px;
}

.app-title {
    font-size: 32px;
    font-weight: 800;
    letter-spacing: -0.04em;
}

.app-subtitle {
    color: #64748b;
    font-size: 15px;
    line-height: 1.5;
    margin-top: 6px;
}

#ocr_output {
    max-height: 640px;
    overflow: auto;
}

.notice {
    margin: 10px auto 16px;
    max-width: 900px;
    padding: 12px 14px;
    border: 1px solid #e5e7eb;
    border-radius: 12px;
    background: #f8fafc;
    color: #334155;
    font-size: 14px;
}
"""

with gr.Blocks(css=custom_css, theme=gr.themes.Soft(), title=APP_TITLE) as demo:
    gr.HTML(
        f"""
        <div class="app-header">
            <div class="app-title">{APP_TITLE}</div>
            <div class="app-subtitle">{APP_SUBTITLE}</div>
        </div>
        """
    )

    gr.HTML(
        f"""
        <div class="notice">
            <strong>Local mode:</strong> OCR runs inside this app using PaddleOCR.
            Current language: <code>{LANG}</code>. PDF render DPI: <code>{PDF_DPI}</code>.
        </div>
        """
    )

    with gr.Row():
        with gr.Column(scale=4):
            upload = gr.File(
                label="Upload image or PDF",
                file_count="single",
                type="filepath",
                file_types=[".pdf", ".png", ".jpg", ".jpeg", ".webp", ".bmp", ".tif", ".tiff"],
            )

            run_btn = gr.Button("Run OCR", variant="primary")

        with gr.Column(scale=8):
            with gr.Tabs():
                with gr.Tab("Text Output"):
                    text_output = gr.Markdown(elem_id="ocr_output")

                with gr.Tab("Structured Output"):
                    json_output = gr.Code(language="json")

    run_btn.click(
        fn=run_local_ocr,
        inputs=[upload],
        outputs=[text_output, json_output],
    )


if __name__ == "__main__":
    demo.queue(max_size=8).launch()