import argparse
import ast
import base64
import hashlib
import json
import os
import re
from html import escape
from io import BytesIO
from pathlib import Path
from typing import Iterable

os.environ.setdefault("HF_MODULES_CACHE", str(Path(os.getenv("TMPDIR", "/tmp")) / "hf_modules_cache"))

import torch
from PIL import Image, ImageFile, ImageOps
from transformers import AutoModelForCausalLM, AutoProcessor

from modeling.modeling_preprocessor import Preprocessor

ImageFile.LOAD_TRUNCATED_IMAGES = True

PROMPTS = {
    "Caption": "Please output the text content from the image.",
    "List-item": "Please output the text content from the image.",
    "Page-footer": "Please output the text content from the image.",
    "Page-header": "Please output the text content from the image.",
    "Section-header": "Please output the text content from the image.",
    "Text": "Please output the text content from the image.",
    "Title": "Please output the text content from the image.",
    "Formula": "Please write out the expression of the formula in the image using LaTeX format.",
    "Table": "Please extract the table from the image and represent it in OTSL format.",
    "Picture": "Please describe the image content.",
    "LAYOUT": "Please output the categories and coordinates of the document elements in reading order.",
}

IMAGE_EXTS = {".jpg", ".jpeg", ".png", ".bmp", ".webp", ".tif", ".tiff"}
RECOGNITION_LABELS = set(PROMPTS) - {"LAYOUT"}


def build_prompt(question: str) -> str:
    return (
        "<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n"
        "<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>"
        f"{question}<|im_end|>\n"
        "<|im_start|>assistant\n"
    )


def make_artifact_filename(stem: str, suffix: str, max_bytes: int = 255) -> str:
    candidate = f"{stem}{suffix}"
    if len(candidate.encode("utf-8")) <= max_bytes:
        return candidate
    digest = hashlib.sha256(candidate.encode("utf-8")).hexdigest()[:10]
    trailer = f"_{digest}{suffix}"
    budget = max_bytes - len(trailer.encode("utf-8"))
    shortened = stem.encode("utf-8")[:budget].decode("utf-8", errors="ignore").rstrip(" .")
    return f"{shortened or 'artifact'}{trailer}"


def image_to_png_data_uri(image: Image.Image) -> str:
    buffer = BytesIO()
    image.convert("RGB").save(buffer, format="PNG")
    encoded = base64.b64encode(buffer.getvalue()).decode("ascii")
    return f"data:image/png;base64,{encoded}"


def save_picture_block(image: Image.Image, image_dir: Path, doc_name: str, sub_idx: int) -> str:
    image_dir.mkdir(parents=True, exist_ok=True)
    image_name = make_artifact_filename(doc_name, f"_sub{sub_idx}.jpg")
    image.convert("RGB").save(image_dir / image_name, format="JPEG", quality=95)
    return f"../images/{image_name}"


def resize_by_pixels(image: Image.Image, max_pixels: int | None) -> Image.Image:
    if not max_pixels or image.width * image.height <= max_pixels:
        return image
    scale = (max_pixels / float(image.width * image.height)) ** 0.5
    size = (max(1, int(image.width * scale)), max(1, int(image.height * scale)))
    return image.resize(size, Image.Resampling.LANCZOS)


def load_image(path: str | Path, max_pixels: int | None = None) -> Image.Image:
    image = Image.open(path)
    image = ImageOps.exif_transpose(image).convert("RGB")
    return resize_by_pixels(image, max_pixels)


def load_pdf_images(path: str | Path, max_pixels: int | None = None) -> list[Image.Image]:
    try:
        import pypdfium2 as pdfium
    except ImportError as exc:
        raise RuntimeError("PDF input requires pypdfium2. Install it or pass image files instead.") from exc

    images = []
    pdf = pdfium.PdfDocument(str(path))
    for page in pdf:
        bitmap = page.render(scale=2.0).to_pil()
        images.append(resize_by_pixels(bitmap.convert("RGB"), max_pixels))
    return images


def iter_documents(input_path: str | Path, max_pixels: int | None) -> Iterable[dict]:
    path = Path(input_path)
    files = [path] if path.is_file() else sorted(p for p in path.rglob("*") if p.is_file())
    for file_path in files:
        suffix = file_path.suffix.lower()
        if suffix in IMAGE_EXTS:
            yield {"name": file_path.stem, "images": [load_image(file_path, max_pixels)]}
        elif suffix == ".pdf":
            yield {"name": file_path.stem, "images": load_pdf_images(file_path, max_pixels)}


def extract_balanced(text: str, left: str, right: str) -> list[str]:
    blocks, depth, start = [], 0, -1
    for i, char in enumerate(text):
        if char == left:
            if depth == 0:
                start = i
            depth += 1
        elif char == right and depth > 0:
            depth -= 1
            if depth == 0 and start != -1:
                blocks.append(text[start:i + 1])
                start = -1
    return list(dict.fromkeys(blocks))


def extract_tolerant_dicts(text: str) -> list[str]:
    blocks = extract_balanced(text, "{", "}")
    for start, char in enumerate(text):
        if char != "{":
            continue
        depth = 0
        end = None
        for cursor in range(start, len(text)):
            if text[cursor] == "{":
                depth += 1
            elif text[cursor] == "}":
                depth -= 1
                if depth == 0:
                    end = cursor + 1
                    break
        blocks.append(text[start:end] if end is not None else text[start:] + ("}" * max(depth, 1)))
    return list(dict.fromkeys(blocks))


def normalize_layout_item(item) -> dict | None:
    if not isinstance(item, dict) or "bbox" not in item or "label" not in item:
        return None
    bbox = item["bbox"]
    if not isinstance(bbox, (list, tuple)) or len(bbox) != 4:
        return None
    try:
        bbox = [float(x) for x in bbox]
    except (TypeError, ValueError):
        return None
    return {"bbox": bbox, "label": str(item["label"])}


def parse_layout_text(text: str) -> list[dict]:
    text = (text or "").strip()
    candidates = [text]
    candidates.extend(extract_balanced(text, "[", "]"))
    first_list = text.find("[")
    if first_list >= 0:
        tail = text[first_list:]
        candidates.append(tail + ("]" * max(0, tail.count("[") - tail.count("]"))))

    best = []
    for candidate in dict.fromkeys(candidates):
        try:
            value = ast.literal_eval(candidate)
        except (SyntaxError, ValueError, TypeError, MemoryError, RecursionError):
            continue
        if not isinstance(value, list):
            continue
        items = [item for raw in value if (item := normalize_layout_item(raw)) is not None]
        if len(items) > len(best):
            best = items
    dict_items = []
    for candidate in extract_tolerant_dicts(text):
        try:
            item = normalize_layout_item(ast.literal_eval(candidate))
        except (SyntaxError, ValueError, TypeError, MemoryError, RecursionError):
            continue
        if item is not None:
            dict_items.append(item)
    if len(dict_items) > len(best):
        best = dict_items
    return best


def map_bbox(bbox: list[float], width: int, height: int) -> list[int]:
    x1, y1, x2, y2 = bbox
    x1, x2 = x1 / 1000.0 * width, x2 / 1000.0 * width
    y1, y2 = y1 / 1000.0 * height, y2 / 1000.0 * height
    if x1 > x2:
        x1, x2 = x2, x1
    if y1 > y2:
        y1, y2 = y2, y1
    x1 = max(0, min(int(round(x1)), max(0, width - 1)))
    y1 = max(0, min(int(round(y1)), max(0, height - 1)))
    x2 = max(x1 + 1, min(int(round(x2)), width))
    y2 = max(y1 + 1, min(int(round(y2)), height))
    return [x1, y1, x2, y2]


def otsl_to_html(otsl: str) -> str:
    if not otsl or not otsl.strip():
        return "<table></table>"

    rows_tokens = otsl.split("<nl>")
    if rows_tokens and rows_tokens[-1] == "":
        rows_tokens.pop()

    grid = []
    for r_idx, row_str in enumerate(rows_tokens):
        if r_idx >= len(grid):
            grid.append([])
        if not row_str.strip():
            continue

        parts = re.findall(r"<([a-z]+)>(.*?)(?=<[a-z]+>|$)", row_str)
        col_idx = 0
        for tag, cell_content in parts:
            while True:
                while len(grid[r_idx]) <= col_idx:
                    grid[r_idx].append(None)
                if grid[r_idx][col_idx] is None:
                    break
                col_idx += 1

            if tag in {"fcel", "ecel"}:
                grid[r_idx][col_idx] = {
                    "text": cell_content.strip() if tag == "fcel" else "",
                    "rowspan": 1,
                    "colspan": 1,
                    "valid": True,
                }
                col_idx += 1
            elif tag == "lcel":
                found = False
                for search_c in range(col_idx - 1, -1, -1):
                    if len(grid[r_idx]) > search_c:
                        cell = grid[r_idx][search_c]
                        if cell and cell.get("valid"):
                            cell["colspan"] += 1
                            found = True
                            break
                grid[r_idx][col_idx] = (
                    {"valid": False, "type": "lcel"}
                    if found
                    else {"text": "", "rowspan": 1, "colspan": 1, "valid": True}
                )
                col_idx += 1
            elif tag == "ucel":
                found = False
                for search_r in range(r_idx - 1, -1, -1):
                    if len(grid[search_r]) > col_idx:
                        cell = grid[search_r][col_idx]
                        if cell and cell.get("valid"):
                            cell["rowspan"] += 1
                            found = True
                            break
                grid[r_idx][col_idx] = (
                    {"valid": False, "type": "ucel"}
                    if found
                    else {"text": "", "rowspan": 1, "colspan": 1, "valid": True}
                )
                col_idx += 1
            elif tag == "xcel":
                grid[r_idx][col_idx] = {"valid": False, "type": "xcel"}
                col_idx += 1
            else:
                col_idx += 1

    html_parts = ["<table>"]
    for row in grid:
        html_parts.append("<tr>")
        for cell in row:
            if cell is None or not cell.get("valid"):
                continue
            attrs = []
            if cell["rowspan"] > 1:
                attrs.append(f'rowspan="{cell["rowspan"]}"')
            if cell["colspan"] > 1:
                attrs.append(f'colspan="{cell["colspan"]}"')
            attr_text = " " + " ".join(attrs) if attrs else ""
            html_parts.append(f"<td{attr_text}>{escape(cell['text'])}</td>")
        html_parts.append("</tr>")
    html_parts.append("</table>")
    return "".join(html_parts)


def process_formula(content: str) -> tuple[str, str | None]:
    content = (content or "").strip("$").strip()
    content = re.sub(r"(?:\\quad\s*){5,}", r"\\quad ", content)
    content = re.sub(r"(?:\\qquad\s*){5,}", r"\\qquad ", content).strip()

    extracted = None
    tag_pattern = (
        r"(?:\\quad|\\qquad|\\eqno)\s*\(([^()]*)\)\s*$"
        r"|\\tag\{([^{}]*)\}\s*$"
    )
    match = re.search(tag_pattern, content)
    if match:
        extracted = match.group(1) or match.group(2)
        content = content[:match.start()].rstrip()

    begin_env = None
    begin_match = re.match(r"^\\begin\{([^}]+)\}", content)
    if begin_match:
        begin_env = begin_match.group(1)
        content = content[begin_match.end():].lstrip()
        end_match = re.search(rf"\\end\{{{re.escape(begin_env)}\}}\s*$", content)
        if end_match:
            content = content[:end_match.start()].rstrip()

    match = re.search(tag_pattern, content)
    if match:
        extracted = match.group(1) or match.group(2)
        content = content[:match.start()].rstrip()

    if begin_env:
        content = f"\\begin{{{begin_env}}}\n{content}\n\\end{{{begin_env}}}"
    return content, extracted


def format_block_content(label: str, raw: str) -> str:
    content = (raw or "").strip()
    if label == "Formula":
        formula, extracted = process_formula(content)
        content = f"$$\n{formula}\n$$"
        if extracted:
            content = f"{content}\n{extracted}"
    elif label == "Table":
        content = content if os.getenv("MOCR2_TABLE_HTML", "0") == "1" else otsl_to_html(content)
    elif label == "Title":
        content = "# " + content.replace("\n", "\n# ")
    elif label == "Section-header":
        content = "## " + content.replace("\n", "\n## ")
    return content


class TransformersMonkeyOCR:
    def __init__(self, model_path: str, device: str | None = None):
        self.device = torch.device(device or ("cuda" if torch.cuda.is_available() else "cpu"))
        dtype = torch.bfloat16 if self.device.type == "cuda" else torch.float32
        self.processor = AutoProcessor.from_pretrained(model_path, trust_remote_code=True, use_fast=True)
        self.processor.tokenizer.padding_side = "left"
        if self.processor.tokenizer.pad_token_id is None:
            self.processor.tokenizer.pad_token = self.processor.tokenizer.eos_token
        self.model = AutoModelForCausalLM.from_pretrained(
            model_path,
            trust_remote_code=True,
            dtype=dtype,
            low_cpu_mem_usage=True,
        ).to(self.device)
        self.model.eval()

    @torch.inference_mode()
    def infer(self, image: Image.Image, question: str, max_new_tokens: int = 4096) -> str:
        return self.batch_infer([image], [question], max_new_tokens=max_new_tokens)[0]

    @torch.inference_mode()
    def batch_infer(
        self,
        images: list[Image.Image],
        questions: list[str],
        max_new_tokens: int = 4096,
        batch_size: int = 4,
    ) -> list[str]:
        if len(images) != len(questions):
            raise ValueError("images and questions must have the same length.")
        if not images:
            return []

        outputs = []
        batch_size = max(1, int(batch_size))
        for start in range(0, len(images), batch_size):
            batch_images = [image.convert("RGB") for image in images[start:start + batch_size]]
            batch_questions = questions[start:start + batch_size]
            outputs.extend(self._generate_batch(batch_images, batch_questions, max_new_tokens))
        return outputs

    def _generate_batch(
        self,
        images: list[Image.Image],
        questions: list[str],
        max_new_tokens: int,
    ) -> list[str]:
        inputs = self.processor(
            text=[build_prompt(question) for question in questions],
            images=images,
            padding=True,
            return_tensors="pt",
        )
        inputs = {
            key: value.to(self.device) if hasattr(value, "to") else value
            for key, value in inputs.items()
        }
        generated = self.model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            do_sample=False,
            temperature=None,
            top_p=None,
        )
        prompt_len = inputs["input_ids"].shape[1]
        output_ids = generated[:, prompt_len:]
        return [
            text.strip()
            for text in self.processor.tokenizer.batch_decode(output_ids, skip_special_tokens=True)
        ]


def parse_page(model: TransformersMonkeyOCR, image: Image.Image) -> tuple[list[dict], str]:
    raw_layout = model.infer(image, PROMPTS["LAYOUT"], max_new_tokens=4096)
    layout = []
    for item in parse_layout_text(raw_layout):
        layout.append({
            "bbox": map_bbox(item["bbox"], image.width, image.height),
            "label": item["label"],
        })

    records = []
    for index, item in enumerate(layout):
        label = item["label"]
        crop = image.crop(item["bbox"])
        if label == "Picture":
            content = ""
        elif label in RECOGNITION_LABELS:
            content = format_block_content(label, model.infer(crop, PROMPTS[label], max_new_tokens=4096))
        else:
            content = ""
        records.append({
            "bbox": item["bbox"],
            "label": label,
            "content": content,
            "block_index": index,
            "_image": crop if label == "Picture" else None,
        })
    return records, raw_layout


def parse_pages(
    model: TransformersMonkeyOCR,
    images: list[Image.Image],
    batch_size: int = 4,
) -> list[dict]:
    raw_layouts = model.batch_infer(
        images,
        [PROMPTS["LAYOUT"]] * len(images),
        max_new_tokens=4096,
        batch_size=batch_size,
    )

    layouts_per_page = []
    for image, raw_layout in zip(images, raw_layouts):
        layout = []
        for item in parse_layout_text(raw_layout):
            layout.append({
                "bbox": map_bbox(item["bbox"], image.width, image.height),
                "label": item["label"],
            })
        layouts_per_page.append(layout)

    crops = []
    questions = []
    owners = []
    for page_idx, (image, layout) in enumerate(zip(images, layouts_per_page)):
        for block_idx, item in enumerate(layout):
            label = item["label"]
            if label != "Picture" and label in RECOGNITION_LABELS:
                crops.append(image.crop(item["bbox"]))
                questions.append(PROMPTS[label])
                owners.append((page_idx, block_idx))

    contents = model.batch_infer(crops, questions, max_new_tokens=4096, batch_size=batch_size)
    page_records = [[] for _ in images]
    for page_idx, (image, layout) in enumerate(zip(images, layouts_per_page)):
        for block_idx, item in enumerate(layout):
            if item["label"] == "Picture":
                page_records[page_idx].append({
                    "bbox": item["bbox"],
                    "label": item["label"],
                    "content": "",
                    "block_index": block_idx,
                    "_image": image.crop(item["bbox"]),
                })
            elif item["label"] not in RECOGNITION_LABELS:
                page_records[page_idx].append({
                    "bbox": item["bbox"],
                    "label": item["label"],
                    "content": "",
                    "block_index": block_idx,
                })

    for content, (page_idx, block_idx) in zip(contents, owners):
        item = layouts_per_page[page_idx][block_idx]
        page_records[page_idx].append({
            "bbox": item["bbox"],
            "label": item["label"],
            "content": format_block_content(item["label"], content),
            "block_index": block_idx,
        })

    for records in page_records:
        records.sort(key=lambda record: record["block_index"])

    return [
        {"records": records, "raw_layout": raw_layout}
        for records, raw_layout in zip(page_records, raw_layouts)
    ]


def prepare_record_for_output(
    record: dict,
    image_dir: Path,
    doc_name: str,
    picture_count: list[int],
    use_base64: bool,
) -> tuple[dict, str]:
    output_record = {key: value for key, value in record.items() if key != "_image"}
    if record.get("label") == "Picture":
        image = record.get("_image")
        if image is not None:
            image_ref = image_to_png_data_uri(image) if use_base64 else save_picture_block(
                image,
                image_dir,
                doc_name,
                picture_count[0],
            )
            picture_count[0] += 1
            output_record["content"] = f"![image]({image_ref})"
    markdown = (output_record.get("content") or "").strip()
    return output_record, markdown


def save_document(
    out_dir: Path,
    doc_name: str,
    page_results: list[dict],
    keep_header_footer: bool = False,
    use_base64: bool = False,
) -> None:
    json_dir = out_dir / "jsons"
    md_dir = out_dir / "markdowns"
    image_dir = out_dir / "images"
    json_dir.mkdir(parents=True, exist_ok=True)
    md_dir.mkdir(parents=True, exist_ok=True)

    records = []
    md_parts = []
    picture_count = [0]
    for page_idx, page in enumerate(page_results):
        if len(page_results) > 1:
            md_parts.append(f"\n\n<!-- page {page_idx + 1} -->\n")
        for record in page["records"]:
            output_record, md = prepare_record_for_output(
                {"page": page_idx + 1, **record},
                image_dir,
                doc_name,
                picture_count,
                use_base64,
            )
            records.append(output_record)
            if md and (keep_header_footer or output_record.get("label") not in {"Page-header", "Page-footer"}):
                md_parts.append(md)

    (json_dir / f"{doc_name}.json").write_text(
        json.dumps(records, ensure_ascii=False, indent=2),
        encoding="utf-8",
    )
    (md_dir / f"{doc_name}.md").write_text("\n\n".join(md_parts).strip() + "\n", encoding="utf-8")


def main():
    parser = argparse.ArgumentParser(description="Minimal Transformers demo for MonkeyOCRv2 two-stage parsing.")
    parser.add_argument("--input-path", "-i", default="../images_test", help="Image/PDF file or directory")
    parser.add_argument("--model-path", "-m", default="../model_weight/MonkeyOCRv2-B-Parsing", help="HF model path")
    parser.add_argument("--output-path", "-o", default="./output/transformers_demo", help="Output directory")
    parser.add_argument("--device", default=None, help="cuda, cuda:0, cpu, ...")
    parser.add_argument("--max-pixels", type=int, default=1003520, help="Resize input pages above this pixel count")
    parser.add_argument("--skip-preprocess", action="store_true", help="Use original pages without preprocessor")
    parser.add_argument("--preprocess-batch-size", type=int, default=8)
    parser.add_argument("--parse-batch-size", type=int, default=8, help="Batch size for Transformers generation")
    parser.add_argument("--keep-header-footer", action="store_true", help="Keep Page-header/Page-footer in markdown")
    parser.add_argument("--use-base64", "--use_base64", action="store_true", help="Embed Picture blocks as base64")
    args = parser.parse_args()

    out_dir = Path(args.output_path)
    parser_model = TransformersMonkeyOCR(args.model_path, device=args.device)
    preprocessor = None
    if not args.skip_preprocess:
        preprocessor = Preprocessor(args.model_path, device=str(parser_model.device), batch_size=args.preprocess_batch_size)

    docs = list(iter_documents(args.input_path, args.max_pixels))
    for doc in docs:
        images = doc["images"]
        if preprocessor is not None:
            images = preprocessor.preprocess_images(images, batch_size=args.preprocess_batch_size)

        for page_idx, image in enumerate(images):
            print(f"Parsing {doc['name']} page {page_idx + 1}/{len(images)}")
        page_results = parse_pages(parser_model, images, batch_size=args.parse_batch_size)
        save_document(
            out_dir,
            doc["name"],
            page_results,
            keep_header_footer=args.keep_header_footer,
            use_base64=args.use_base64,
        )

    print(f"Done. Results saved to {out_dir}")


if __name__ == "__main__":
    main()