#!/usr/bin/env python3
import os
import sys
import argparse
from pathlib import Path
import pypdfium2 as pdfium
import re
from PIL import Image

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)) + "/src")
from booktrans.agent import make_agent, AgentError, RateLimited


SYSTEM_PROMPT = """You are an expert document OCR and layout extraction model.
Extract all text, math, tables, footnotes, and image captions from the provided document page.
Output the result in standard Markdown format.

CRITICAL RULES:
1. TABLES: Convert tables into proper Markdown tables.
2. MULTI-COLUMN: If the page has 2, 3, or 4 columns, read them in the correct reading order (top-to-bottom, left-to-right). Do not mix text from different columns into the same paragraph.
3. MATH: Wrap all inline mathematical formulas in single dollar signs: `$formula$`. Wrap display equations (standalone lines) in double dollar signs: `$$formula$$`. Use valid LaTeX syntax.
4. FOOTNOTES: Format footnotes exactly as `[^1]` in the text and place the footnote content at the bottom of the output as `[^1]: Note text`.
5. MARGINALIA & SIDENOTES: Integrate marginalia and sidenotes into the text flow where they logically belong, or place them at the end of the section.
6. IMAGES & GRAPHS: For every image, graph, diagram, or chart, insert a tag `![image]([ymin, xmin, ymax, xmax])`, where coordinates are integers from 0 to 1000 representing the bounding box normalized to the page size.
7. CAPTIONS: If an image has a caption (legend), extract it and place it immediately after the image tag in italics: `*Caption text*`.
8. HEADERS & FOOTERS: DO NOT extract running headers, running footers, or page numbers. Stitch sentences across pages seamlessly if necessary.
9. NO CODE BLOCKS: Do NOT wrap your output in ```markdown code blocks. Return the raw Markdown directly.
10. HEADINGS: Preserve document hierarchy by using Markdown headings (`#`, `##`, `###`) for section titles and chapters based on their visual prominence (font size, weight).
"""

def extract_page(pdf_path, page_num, total_pages, agent, work_dir):
    page_md_file = work_dir / f"page_{page_num:04d}.md"
    if page_md_file.exists():
        return page_md_file.read_text(encoding="utf-8")
    
    pdf = pdfium.PdfDocument(str(pdf_path))
    page = pdf[page_num - 1]
    bitmap = page.render(scale=3)
    pil_image = bitmap.to_pil()
    
    img_path = work_dir / f"page_{page_num:04d}.png"
    pil_image.save(img_path, format="PNG")
    
    user_prompt = f"Extract page {page_num} of {total_pages}."
    print(f"[*] Processing page {page_num}/{total_pages} with {agent.kind}...")
    
    while True:
        try:
            text, stats = agent.run(SYSTEM_PROMPT, user_prompt, image=str(img_path))
            
            cost = stats.get('cost_usd') or 0.0
            print(f"    Done! Model: {stats.get('model')}, Cost: ${cost:.4f}")
            
            text = text.strip()
            if text.startswith("```markdown"): text = text[11:]
            if text.startswith("```"): text = text[3:]
            if text.endswith("```"): text = text[:-3]
            text = text.strip()
            
            # Crop images based on coordinates
            images_dir = work_dir / "images"
            images_dir.mkdir(parents=True, exist_ok=True)
            
            def replace_img(match):
                caption = match.group(1)
                coords_str = match.group(2)
                coords = [int(c.strip()) for c in coords_str.split(",") if c.strip().isdigit()]
                if len(coords) == 4:
                    ymin, xmin, ymax, xmax = coords
                    W, H = pil_image.size
                    pad_w = int(W * 0.004)
                    pad_h = int(H * 0.004)
                    c_left = max(0, int(xmin * W / 1000.0) - pad_w)
                    c_top = max(0, int(ymin * H / 1000.0) - pad_h)
                    c_right = min(W, int(xmax * W / 1000.0) + pad_w)
                    c_bottom = min(H, int(ymax * H / 1000.0) + pad_h)
                    
                    if c_right > c_left and c_bottom > c_top:
                        cropped = pil_image.crop((c_left, c_top, c_right, c_bottom))
                        img_filename = f"img_p{page_num:04d}_{c_top}_{c_left}.png"
                        cropped.save(images_dir / img_filename, format="PNG")
                        return f"![{caption}](images/{img_filename})"
                return match.group(0)
            
            text = re.sub(r"!\[(.*?)\]\(\[([^\]]+)\]\)", replace_img, text)
            
            page_md_file.write_text(text, encoding="utf-8")
            return text
        except RateLimited:
            print("    Rate limited, waiting 10s...")
            import time; time.sleep(10)
        except AgentError as e:
            print(f"    Agent error: {e}")
            raise

def main():
    parser = argparse.ArgumentParser(description="Universal PDF to Markdown converter using codex or agy.")
    parser.add_argument("pdf", help="Path to the PDF file")
    parser.add_argument("--agent", choices=["codex", "agy"], default="codex", help="Agent to use (codex or agy)")
    parser.add_argument("--model", help="Specific model to use (e.g. gpt-4o)")
    parser.add_argument("--effort", help="Reasoning effort if supported")
    parser.add_argument("--pages", help="Comma-separated list of page numbers to extract (e.g., 5,6,10)")
    args = parser.parse_args()

    pdf_path = Path(args.pdf)
    if not pdf_path.exists():
        sys.exit(f"File not found: {pdf_path}")
        
    work_dir = pdf_path.with_suffix('.work') / 'pdf_pages'
    work_dir.mkdir(parents=True, exist_ok=True)
    
    agent = make_agent(kind=args.agent, model=args.model, effort=args.effort, timeout=300)
    
    pdf = pdfium.PdfDocument(str(pdf_path))
    total_pages = len(pdf)
    
    pages_to_extract = range(1, total_pages + 1)
    if args.pages:
        pages_to_extract = [int(p.strip()) for p in args.pages.split(",")]
    
    final_md_path = pdf_path.with_suffix('.md')
    print(f"Starting extraction of {len(pages_to_extract)} pages into {final_md_path}...")
    
    all_text = []
    for i in pages_to_extract:
        text = extract_page(pdf_path, i, total_pages, agent, work_dir)
        all_text.append(text)
        
    final_md_path.write_text("\n\n---\n\n".join(all_text), encoding="utf-8")
    print(f"Extraction complete! Saved to {final_md_path}")
    print(f"You can now run `./bt_codex {final_md_path}`")

if __name__ == "__main__":
    main()
