import gc
import importlib.util
import inspect
import re
import tempfile
import time
from io import BytesIO
from pathlib import Path
from threading import Lock, Thread

import gradio as gr
import spaces
import torch
import transformers
from PIL import Image, ImageOps

try:
    import fitz
except ImportError:
    fitz = None

from transformers import (
    AutoModelForImageTextToText,
    AutoProcessor,
    TextIteratorStreamer,
)

MAX_MAX_NEW_TOKENS = 8192
DEFAULT_MAX_NEW_TOKENS = 4096

MODEL_PATH = "zai-org/GLM-OCR"
MIN_TRANSFORMERS_VERSION = "5.0.0"
UPGRADE_TRANSFORMERS_CMD = 'pip install -U "transformers>=5.0.0"'
MODEL_CARD_TRANSFORMERS_CMD = "pip install git+https://github.com/huggingface/transformers.git"

BASE_DIR = Path(__file__).resolve().parent
DATA_DIR = BASE_DIR / "data"

INPUT_MODES = ["Upload PDF", "PDF from data", "All PDFs in data"]
DEFAULT_INPUT_MODE = "Upload PDF"

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print("Using device:", device)

processor = None
model = None
model_load_error = None
model_init_lock = Lock()


def parse_version_triplet(version_text):
    parts = [int(part) for part in re.findall(r"\d+", str(version_text))[:3]]
    while len(parts) < 3:
        parts.append(0)
    return tuple(parts)


def build_glm_ocr_dependency_error():
    transformers_version = getattr(transformers, "__version__", "unknown")
    try:
        import torchvision  # noqa: F401
    except ImportError:
        torch_version = getattr(torch, "__version__", "unknown")
        return (
            "GLM-OCR requires torchvision, but it is not installed in the active environment. "
            f"Install a torchvision build compatible with torch {torch_version} "
            "(for example, `pip install torchvision`) and restart the app."
        )
    except Exception as exc:
        return (
            "GLM-OCR could not import torchvision from the active environment. "
            f"Resolve the torchvision installation issue ({exc}) and restart the app."
        )

    has_glm_ocr_support = importlib.util.find_spec("transformers.models.glm_ocr") is not None
    if has_glm_ocr_support:
        return None

    version_triplet = parse_version_triplet(transformers_version)
    minimum_triplet = parse_version_triplet(MIN_TRANSFORMERS_VERSION)
    if version_triplet and version_triplet < minimum_triplet:
        return (
            f"GLM-OCR support is not available in transformers {transformers_version}. "
            f"Upgrade to {MIN_TRANSFORMERS_VERSION}+ with `{UPGRADE_TRANSFORMERS_CMD}` or use the "
            f"model card recommendation `{MODEL_CARD_TRANSFORMERS_CMD}`, then restart the app."
        )

    return (
        f"GLM-OCR support is not available in the installed transformers build ({transformers_version}). "
        f"Upgrade transformers with `{UPGRADE_TRANSFORMERS_CMD}` or use the model card recommendation "
        f"`{MODEL_CARD_TRANSFORMERS_CMD}`, then restart the app."
    )


def load_glm_ocr_components():
    global processor, model, model_load_error

    if model_load_error is not None:
        raise RuntimeError(model_load_error)

    if processor is not None and model is not None:
        return processor, model

    with model_init_lock:
        if model_load_error is not None:
            raise RuntimeError(model_load_error)

        if processor is not None and model is not None:
            return processor, model

        dependency_error = build_glm_ocr_dependency_error()
        if dependency_error is not None:
            model_load_error = dependency_error
            print(model_load_error)
            raise RuntimeError(model_load_error)

        try:
            processor = AutoProcessor.from_pretrained(MODEL_PATH, trust_remote_code=True)

            model_kwargs = {
                "pretrained_model_name_or_path": MODEL_PATH,
                "torch_dtype": torch.bfloat16 if torch.cuda.is_available() else torch.float32,
                "trust_remote_code": True,
            }
            if torch.cuda.is_available():
                model_kwargs["device_map"] = "auto"

            model = AutoModelForImageTextToText.from_pretrained(**model_kwargs).eval()
            if not torch.cuda.is_available():
                model = model.to(device)
        except Exception as exc:
            model_load_error = (
                f"Failed to load {MODEL_PATH}: {exc}. "
                "If the error mentions an unrecognized processor or model type, upgrade transformers with "
                f"`{UPGRADE_TRANSFORMERS_CMD}` or follow the model card recommendation "
                f"`{MODEL_CARD_TRANSFORMERS_CMD}`. Restart the app after upgrading."
            )
            print(model_load_error)
            raise RuntimeError(model_load_error) from exc

    return processor, model


TASK_PROMPTS = {
    "Text": "Text Recognition:",
    "Formula": "Formula Recognition:",
    "Table": "Table Recognition:",
}

TASK_CHOICES = list(TASK_PROMPTS.keys())


def list_data_pdfs():
    if not DATA_DIR.exists():
        return []

    return sorted(
        [path for path in DATA_DIR.iterdir() if path.is_file() and path.suffix.lower() == ".pdf"],
        key=lambda path: path.name.lower(),
    )


def build_data_folder_note():
    pdf_paths = list_data_pdfs()
    folder_line = f"Data folder: `{DATA_DIR}`"
    if not pdf_paths:
        return folder_line + "\n\nNo PDF files were found."

    lines = "\n".join(f"- `{path.name}`" for path in pdf_paths)
    return folder_line + "\n\nAvailable PDFs:\n" + lines


def refresh_pdf_dropdown():
    pdf_names = [path.name for path in list_data_pdfs()]
    value = pdf_names[0] if pdf_names else None
    return gr.update(choices=pdf_names, value=value), build_data_folder_note()


def resolve_data_pdf_path(pdf_name):
    if not pdf_name:
        raise ValueError("Please choose a PDF from the data folder.")

    for path in list_data_pdfs():
        if path.name == pdf_name:
            return path

    raise ValueError(f"Could not find `{pdf_name}` in `{DATA_DIR}`.")


def resolve_uploaded_pdf_path(uploaded_pdf):
    if not uploaded_pdf:
        raise ValueError("Please upload a PDF first.")

    pdf_path = Path(str(uploaded_pdf))
    if not pdf_path.exists():
        raise ValueError("The uploaded PDF could not be read. Please upload it again.")

    if pdf_path.suffix.lower() != ".pdf":
        raise ValueError("Please upload a PDF file.")

    return pdf_path


def cache_uploaded_pdf(uploaded_pdf):
    if not uploaded_pdf:
        return None, "No PDF uploaded yet."

    pdf_path = resolve_uploaded_pdf_path(uploaded_pdf)
    return str(pdf_path), f"Selected upload: `{pdf_path.name}`"


def normalize_image(image: Image.Image):
    if image.mode in ("RGBA", "LA", "P"):
        image = image.convert("RGB")
    return ImageOps.exif_transpose(image)


def pdf_page_to_image(pdf_path, page_number):
    if fitz is None:
        raise RuntimeError("PDF support requires PyMuPDF (`fitz`) to be installed.")

    page_number = int(page_number)
    if page_number < 1:
        raise ValueError("Page number must be 1 or higher.")

    with fitz.open(pdf_path) as document:
        if page_number > document.page_count:
            raise ValueError(
                f"Page {page_number} is outside the page count for {Path(pdf_path).name} "
                f"({document.page_count} pages)."
            )
        page = document.load_page(page_number - 1)
        pixmap = page.get_pixmap(matrix=fitz.Matrix(2, 2), alpha=False)
        image = Image.open(BytesIO(pixmap.tobytes("png"))).convert("RGB")
        return ImageOps.exif_transpose(image)


def get_pdf_page_count(pdf_path):
    if fitz is None:
        raise RuntimeError("PDF support requires PyMuPDF (`fitz`) to be installed.")

    with fitz.open(pdf_path) as document:
        return document.page_count


def parse_page_selection(page_selection, page_count):
    selection = str(page_selection or "1").strip().lower()
    if not selection:
        selection = "1"
    if selection == "all":
        return list(range(1, page_count + 1))

    pages = []
    for raw_part in selection.split(","):
        part = raw_part.strip()
        if not part:
            continue
        if "-" in part:
            start_str, end_str = [token.strip() for token in part.split("-", 1)]
            start_page = int(start_str)
            end_page = int(end_str)
            if start_page > end_page:
                raise ValueError(f"Invalid page range: {part}")
            pages.extend(range(start_page, end_page + 1))
        else:
            pages.append(int(part))

    if not pages:
        raise ValueError("Please provide page numbers such as `1`, `1,3`, `2-5`, or `all`.")

    deduped_pages = []
    seen = set()
    for page in pages:
        if page < 1 or page > page_count:
            raise ValueError(f"Page {page} is outside the PDF page count ({page_count}).")
        if page not in seen:
            deduped_pages.append(page)
            seen.add(page)
    return deduped_pages


def join_output_blocks(blocks):
    return "\n\n".join(block.strip() for block in blocks if str(block).strip()).strip()


def stream_pdf_text(pdf_path, task, page_selection, max_new_tokens):
    page_count = get_pdf_page_count(pdf_path)
    selected_pages = parse_page_selection(page_selection, page_count)
    include_page_headers = len(selected_pages) > 1
    combined_sections = []

    for page_number in selected_pages:
        page_title = f"Page {page_number}"
        try:
            image = pdf_page_to_image(pdf_path, page_number)
            page_text = ""
            for chunk in process_image_stream(
                image=image,
                task=task,
                max_new_tokens=max_new_tokens,
            ):
                page_text = chunk.strip()
                current_block = f"{page_title}\n{page_text}" if include_page_headers else page_text
                yield join_output_blocks(combined_sections + [current_block])

            if not page_text.strip():
                page_text = "[ERROR] No output was generated."
        except Exception as exc:
            page_text = f"[ERROR] {str(exc)}"

        final_block = f"{page_title}\n{page_text}" if include_page_headers else page_text
        combined_sections.append(final_block)
        yield join_output_blocks(combined_sections)


def stream_data_folder_pdfs(task, page_selection, max_new_tokens):
    pdf_paths = list_data_pdfs()
    if not pdf_paths:
        yield f"[ERROR] No PDF files were found in `{DATA_DIR}`."
        return

    combined_files = []
    for path in pdf_paths:
        pdf_text = ""
        try:
            for partial in stream_pdf_text(path, task, page_selection, max_new_tokens):
                pdf_text = partial
                yield join_output_blocks(combined_files + [f"File: {path.name}\n{pdf_text}"])

            if not pdf_text.strip():
                pdf_text = "[ERROR] No output was generated."
            file_block = f"File: {path.name}\n{pdf_text}"
        except Exception as exc:
            file_block = f"File: {path.name}\n[ERROR] {str(exc)}"
            yield join_output_blocks(combined_files + [file_block])

        combined_files.append(file_block)
        yield join_output_blocks(combined_files)


def calc_timeout_generic(*args, **kwargs):
    gpu_timeout = kwargs.get("gpu_timeout", None)
    if gpu_timeout is None and args:
        gpu_timeout = args[-1]
    try:
        return int(gpu_timeout)
    except Exception:
        return 60


def process_image_stream(image, task, max_new_tokens=DEFAULT_MAX_NEW_TOKENS, gpu_timeout=60):
    del gpu_timeout
    tmp_path = None
    try:
        if image is None:
            yield "[ERROR] Please upload an image first."
            return

        if task not in TASK_PROMPTS:
            yield "[ERROR] Invalid OCR task selected."
            return

        processor_obj, model_obj = load_glm_ocr_components()
        image = normalize_image(image)

        tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".png")
        image.save(tmp.name, "PNG")
        tmp_path = tmp.name
        tmp.close()

        prompt = TASK_PROMPTS[task]
        messages = [
            {
                "role": "user",
                "content": [
                    {"type": "image", "url": tmp_path},
                    {"type": "text", "text": prompt},
                ],
            }
        ]

        inputs = processor_obj.apply_chat_template(
            messages,
            tokenize=True,
            add_generation_prompt=True,
            return_dict=True,
            return_tensors="pt",
        )

        inputs.pop("token_type_ids", None)
        inputs = {key: value.to(model_obj.device) if hasattr(value, "to") else value for key, value in inputs.items()}

        streamer = TextIteratorStreamer(
            processor_obj.tokenizer if hasattr(processor_obj, "tokenizer") else processor_obj,
            skip_prompt=True,
            skip_special_tokens=True,
        )

        generation_error = {"error": None}
        generation_kwargs = {
            **inputs,
            "streamer": streamer,
            "max_new_tokens": int(max_new_tokens),
        }

        def run_generation():
            try:
                model_obj.generate(**generation_kwargs)
            except Exception as exc:
                generation_error["error"] = exc
                try:
                    streamer.end()
                except Exception:
                    pass

        thread = Thread(target=run_generation, daemon=True)
        thread.start()

        buffer = ""
        for new_text in streamer:
            buffer += new_text
            time.sleep(0.01)
            yield buffer.strip()

        thread.join(timeout=1.0)

        if generation_error["error"] is not None:
            error_message = f"[ERROR] Inference failed: {generation_error['error']}"
            if buffer.strip():
                yield buffer.strip() + "\n\n" + error_message
            else:
                yield error_message
            return

        if not buffer.strip():
            yield "[ERROR] No output was generated."
    except Exception as exc:
        yield f"[ERROR] {str(exc)}"
    finally:
        if tmp_path and Path(tmp_path).exists():
            try:
                Path(tmp_path).unlink()
            except Exception:
                pass
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()


def toggle_input_mode(mode):
    mode = str(mode or DEFAULT_INPUT_MODE)
    if mode == "Upload PDF":
        help_text = "Click the upload button to choose a PDF, then enter pages like `1`, `1,3`, `2-5`, or `all`."
        return (
            gr.update(visible=True),
            gr.update(visible=True),
            gr.update(visible=False),
            gr.update(visible=True),
            gr.update(visible=False),
            gr.update(visible=False),
            gr.update(value=help_text),
        )

    if mode == "All PDFs in data":
        help_text = (
            "Run OCR on every PDF in the data folder. Page selections such as `1`, `1,3`, `2-5`, "
            "or `all` are applied to each file."
        )
        return (
            gr.update(visible=False),
            gr.update(visible=False),
            gr.update(visible=False),
            gr.update(visible=True),
            gr.update(visible=True),
            gr.update(visible=True),
            gr.update(value=help_text),
        )

    help_text = (
        "Choose one PDF from the data folder and enter pages like `1`, `1,3`, `2-5`, or `all`."
    )
    return (
        gr.update(visible=False),
        gr.update(visible=False),
        gr.update(visible=True),
        gr.update(visible=True),
        gr.update(visible=True),
        gr.update(visible=True),
        gr.update(value=help_text),
    )


@spaces.GPU(duration=calc_timeout_generic)
def run_router(input_mode, task, uploaded_pdf, pdf_name, page_selection, max_new_tokens_v, gpu_timeout_v):
    try:
        mode = str(input_mode or DEFAULT_INPUT_MODE)
        if mode == "All PDFs in data":
            yield from stream_data_folder_pdfs(
                task=task,
                page_selection=page_selection,
                max_new_tokens=max_new_tokens_v,
            )
            return

        if mode == "Upload PDF":
            pdf_path = resolve_uploaded_pdf_path(uploaded_pdf)
        else:
            pdf_path = resolve_data_pdf_path(pdf_name)

        yield from stream_pdf_text(
            pdf_path=pdf_path,
            task=task,
            page_selection=page_selection,
            max_new_tokens=max_new_tokens_v,
        )
    except Exception as exc:
        yield f"[ERROR] {str(exc)}"


available_pdfs = [path.name for path in list_data_pdfs()]
default_pdf = available_pdfs[0] if available_pdfs else None

with gr.Blocks(title="GLM-OCR") as demo:
    gr.Markdown("# GLM-OCR")
    gr.Markdown("Upload a PDF with the button below or run OCR on PDFs stored in `glmocr/data`.")

    with gr.Row():
        input_mode = gr.Radio(
            choices=INPUT_MODES,
            value=DEFAULT_INPUT_MODE,
            label="Input mode",
        )
        task = gr.Radio(
            choices=TASK_CHOICES,
            value="Text",
            label="OCR task",
        )

    input_help = gr.Markdown(
        "Click the upload button to choose a PDF, then enter pages like `1`, `1,3`, `2-5`, or `all`."
    )

    with gr.Column():
        uploaded_pdf_state = gr.State(None)
        upload_pdf_btn = gr.UploadButton(
            "Upload PDF",
            file_types=[".pdf"],
            file_count="single",
            type="filepath",
            visible=True,
        )
        uploaded_pdf_status = gr.Markdown("No PDF uploaded yet.", visible=True)
        pdf_dropdown = gr.Dropdown(
            choices=available_pdfs,
            value=default_pdf,
            label="PDF from data folder",
            visible=False,
        )
        page_selection = gr.Textbox(
            value="1",
            label="Pages",
            placeholder="Examples: 1, 1,3, 2-5, all",
            visible=True,
        )
        data_folder_note = gr.Markdown(build_data_folder_note(), visible=False)
        refresh_pdfs_btn = gr.Button("Refresh PDF list", visible=False)

    with gr.Row():
        max_new_tokens = gr.Slider(
            minimum=1,
            maximum=MAX_MAX_NEW_TOKENS,
            step=1,
            value=DEFAULT_MAX_NEW_TOKENS,
            label="Max new tokens",
        )
        gpu_duration_state = gr.Number(
            value=60,
            precision=0,
            label="GPU duration (seconds)",
        )

    run_btn = gr.Button("Run OCR", variant="primary")
    result = gr.Textbox(label="OCR output", lines=24, max_lines=40)

    demo.load(
        fn=refresh_pdf_dropdown,
        inputs=None,
        outputs=[pdf_dropdown, data_folder_note],
        queue=False,
    )

    input_mode.change(
        fn=toggle_input_mode,
        inputs=[input_mode],
        outputs=[
            upload_pdf_btn,
            uploaded_pdf_status,
            pdf_dropdown,
            page_selection,
            data_folder_note,
            refresh_pdfs_btn,
            input_help,
        ],
        queue=False,
    )

    upload_pdf_btn.upload(
        fn=cache_uploaded_pdf,
        inputs=[upload_pdf_btn],
        outputs=[uploaded_pdf_state, uploaded_pdf_status],
        queue=False,
    )

    refresh_pdfs_btn.click(
        fn=refresh_pdf_dropdown,
        inputs=None,
        outputs=[pdf_dropdown, data_folder_note],
        queue=False,
    )

    run_btn.click(
        fn=run_router,
        inputs=[
            input_mode,
            task,
            uploaded_pdf_state,
            pdf_dropdown,
            page_selection,
            max_new_tokens,
            gpu_duration_state,
        ],
        outputs=[result],
    )


if __name__ == "__main__":
    launch_signature = inspect.signature(demo.launch)
    launch_kwargs = {
        "show_error": True,
    }

    if "ssr_mode" in launch_signature.parameters:
        launch_kwargs["ssr_mode"] = False

    if "mcp_server" in launch_signature.parameters:
        launch_kwargs["mcp_server"] = True

    demo.queue(max_size=50).launch(**launch_kwargs)