"""
╔══════════════════════════════════════════════════════════════════╗
║         PDF → Markdown OCR  |  2-Step Pipeline                  ║
║                                                                  ║
║  STEP 1 — EXTRACT  (Vision model: Llama 4 Scout / Maverick)      ║
║    → Faithful, exact OCR of every element on the page            ║
║                                                                  ║
║  STEP 2 — REWRITE  (Text model: Kimi K2 / Llama 70B / Qwen 32B) ║
║    → Creative rewrite + beautiful inline SVG diagrams            ║
║                                                                  ║
║  Features:                                                       ║
║    • Chapter/Unit auto-detection → separate folders              ║
║    • Figures → inline SVG (renders in Typora/Obsidian)           ║
║    • Mermaid for flowcharts and graphs                           ║
║    • Multi-language output                                       ║
║    • Resume support, rate-limit handling                         ║
╚══════════════════════════════════════════════════════════════════╝
"""

import os, base64, argparse, time, re, json
from pathlib import Path
from groq import Groq
import fitz  # PyMuPDF


# ══════════════════════════════════════════════════════════════════════════════
#  MODEL REGISTRY
# ══════════════════════════════════════════════════════════════════════════════

VISION_MODELS = {
    "scout":    "meta-llama/llama-4-scout-17b-16e-instruct",
    "maverick": "meta-llama/llama-4-maverick-17b-128e-instruct",
}

REWRITE_MODELS = {
    "llama70": "llama-3.1-8b-instant",
    "qwen32":  "qwen/qwen3-32b",
    "scout":   "meta-llama/llama-4-scout-17b-16e-instruct",
    "kimi":    "moonshotai/kimi-k2-instruct-0905",
}

DEFAULT_VISION_MODEL  = "scout"
DEFAULT_REWRITE_MODEL = "kimi"


# ══════════════════════════════════════════════════════════════════════════════
#  SUPPORTED LANGUAGES
# ══════════════════════════════════════════════════════════════════════════════

SUPPORTED_LANGUAGES = {
    "en": "English",    "hi": "Hindi",     "mr": "Marathi",   "gu": "Gujarati",
    "bn": "Bengali",    "ta": "Tamil",     "te": "Telugu",    "kn": "Kannada",
    "ml": "Malayalam",  "pa": "Punjabi",   "ur": "Urdu",      "fr": "French",
    "de": "German",     "es": "Spanish",   "ja": "Japanese",  "zh": "Chinese (Simplified)",
    "ar": "Arabic",     "ru": "Russian",   "pt": "Portuguese","ko": "Korean",
}


# ══════════════════════════════════════════════════════════════════════════════
#  IMAGE HELPERS
# ══════════════════════════════════════════════════════════════════════════════

def pdf_page_to_base64(page, dpi: int = 150) -> str:
    pix = page.get_pixmap(dpi=dpi)
    return base64.b64encode(pix.tobytes("jpeg")).decode("utf-8")


# ══════════════════════════════════════════════════════════════════════════════
#  POST-PROCESSING
# ══════════════════════════════════════════════════════════════════════════════

def fix_table_row_math(row: str) -> str:
    row = re.sub(r"\\\((.+?)\\\)", r"$\1$", row)
    row = re.sub(r"\\\[(.+?)\\\]", r"$\1$", row)
    def protect_pipes(m):
        inner = m.group(1).replace('|', r'\\vert ')
        return f"${inner}$"
    row = re.sub(r"\$(.+?)\$", protect_pipes, row, flags=re.DOTALL)
    return re.sub(r"  +", " ", row)

def fix_math_in_tables(content: str) -> str:
    fixed = []
    for line in content.split("\n"):
        if re.match(r"^\s*\|", line):
            fixed.append(fix_table_row_math(line))
        else:
            line = re.sub(r"\\\\\((.+?)\\\\\)", r"$\1$", line)
            line = re.sub(r"\\\[(.+?)\\\]", r"$$\1$$", line, flags=re.DOTALL)
            fixed.append(line)
    return "\n".join(fixed)

def clean_broken_image_links(content: str) -> str:
    content = re.sub(
        r"^!\[.*?\]\(https?://(?:example\.com|placeholder)/.*?\)\s*$",
        "", content, flags=re.MULTILINE
    )
    content = re.sub(r"!\[.*?\]\(https?://(?:example\.com|placeholder)/.*?\)", "", content)
    content = re.sub(r"!\[.*?\]\((?:\./)?\w+\.(?:png|jpg|jpeg|svg|gif)\)", "", content)
    return content

def post_process(content: str) -> str:
    content = fix_math_in_tables(content)
    content = clean_broken_image_links(content)
    return content


# ══════════════════════════════════════════════════════════════════════════════
#  CHAPTER DETECTION
# ══════════════════════════════════════════════════════════════════════════════

CHAPTER_DETECTION_PROMPT = """
******************STRICTLY FOLLOW THIS ****************
Analyze this academic book page image.
Respond ONLY with a JSON object — no other text whatsoever.

Decide: is this page the START of a brand-new chapter / unit / major division?

═════════════════════════════════════════════════════
TRUE — set is_new_chapter: true — only when ALL hold:
═════════════════════════════════════════════════════
1. A LARGE PROMINENT heading is the dominant text on the page — clearly introduces
   a NEW major division of the book (not a sub-section or a running header).


***********MOST IMPORTANT DONOT MISS*****************
2. The heading uses one of these division labels (any case, any language script):
   English  → Chapter, Unit, Part, Module, Lesson, Topic, Section, Volume, Block, Theme
   Nepali   → अध्याय, एकाइ, भाग
   Hindi    → अध्याय, इकाई, भाग
   Bengali  → অধ্যায়
   Punjabi  → ਅਧਿਆਇ
   Tamil    → அத்தியாயம்
   (or any recognisable equivalent in other scripts)

3. The heading includes a NUMBER (Arabic 1,2,3 OR Roman I,II,III OR spelled-out One,Two...)
   PLUS a descriptive title. Accepted formats include:
     "Chapter 1 – Rotational Dynamics"
     "UNIT III: Thermodynamics"   /   "Unit 3"  (title may be on the next line)
     "Part Two — Organic Chemistry"
     "Module 7: Data Structures"
     "Lesson IV: The French Revolution"
     "अध्याय ३ – गतिकी"
     A standalone large number is ALSO acceptable when accompanied by a title
     directly below it (many textbooks format it this way).

4. Layout is consistent with a division opener: large whitespace above body text,
   an illustration/graphic below the heading, or body text simply begins on this page.

═════════════════════════════════════════════════════
FALSE — set is_new_chapter: false — for ANY of these:
═════════════════════════════════════════════════════
- Mid-chapter page (body text continues from a previous page with no major heading)
- Sub-section only: e.g. "1.1 Angular Displacement", "3.2 Energy", "Section 4.3"
- Table of Contents, Index, Preface, Foreword, Acknowledgements, Bibliography, Glossary
- Appendix, Answer Key, Formula Sheet, Question Bank, Practice Problems page
- Chapter name visible ONLY in a running header/footer at top or bottom edge,
  not as the main large heading on the page
- Continuation of a chapter already-started, even if the chapter title is printed at top
- Blank or nearly blank page

════════════════════════════════
 Field extraction rules:
════════════════════════════════
chapter_number : number portion ONLY → "1", "2", "III", "IV", "A"
                 Convert spelled-out to digit ("Chapter One" → "1")
                 null if absent
chapter_title  : descriptive title ONLY, WITHOUT the label+number prefix
                 ("Chapter 1 Rotational Dynamics" → "Rotational Dynamics")
                 If the title is on a separate line below the number, include it.
                 null if false
chapter_type   : exact label word found, first-letter capitalised
                 e.g. "Chapter", "Unit", "Part", "Module", "Lesson"
                 Default to "Chapter" when unclear
confidence     : integer 0–100 (your certainty about is_new_chapter)

Return EXACTLY one JSON object, nothing else:
{"is_new_chapter": true,  "chapter_number": "1",  "chapter_title": "Rotational Dynamics", "chapter_type": "Chapter", "confidence": 97}
OR
{"is_new_chapter": false, "chapter_number": null, "chapter_title": null, "chapter_type": null, "confidence": 92}
"""

def detect_chapter(client: Groq, base64_image: str, vision_model: str) -> dict:
    try:
        response = client.chat.completions.create(
            model=vision_model,
            messages=[{"role": "user", "content": [
                {"type": "text",      "text": CHAPTER_DETECTION_PROMPT},
                {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{base64_image}"}},
            ]}],
            temperature=0.0, max_completion_tokens=300,
        )
        raw = re.sub(r"^```(?:json)?\s*|\s*```$", "", response.choices[0].message.content.strip())
        result = json.loads(raw)
        conf = result.get("confidence", "?")
        if result.get("is_new_chapter"):
            print(f"[conf={conf}%]", end=" ", flush=True)
        return result
    except Exception as e:
        print(f"\n    ⚠️  Chapter detection failed ({e})")
        return {"is_new_chapter": False, "chapter_number": None, "chapter_title": None, "chapter_type": None, "confidence": 0}

def slugify(text: str) -> str:
    text = re.sub(r"[^\w\s-]", "", text.lower().strip())
    return re.sub(r"[\s_-]+", "_", text)[:60]

def make_chapter_folder(ctype, cnum, ctitle) -> str:
    ctype  = (ctype or "chapter").lower()
    cnum   = str(cnum or "00").zfill(2) if re.match(r"^\d+$", str(cnum or "00")) else (cnum or "00")
    return f"{ctype}_{cnum}_{slugify(ctitle or 'untitled')}"


# ══════════════════════════════════════════════════════════════════════════════
#  STEP 1 — EXTRACTION  (faithful OCR)
# ══════════════════════════════════════════════════════════════════════════════

EXTRACTION_PROMPT = """You are a precise OCR engine. Your ONLY job is to faithfully extract content.

⛔ ABSOLUTE PROHIBITIONS — violating ANY of these is a critical failure:
- NO step-by-step reasoning ("Step 1:", "Step 2:", "Step 3:" etc.)
- NO meta-commentary ("The content includes...", "This page contains...", "I will extract...")
- NO chain-of-thought or self-explanation
- NO "The final answer is:", "In summary:", or any closing remarks
- NO repeating the same content twice
- NO describing what you are about to do — just DO it
- Start IMMEDIATELY with the actual page content. First character must be content.

OUTPUT FORMAT (content only, no wrapper):
- Markdown for structure (headings, lists, tables)
- $...$ for inline LaTeX, $$...$$ for block LaTeX
- Table cells: LaTeX on single line, use \\vert not |
- For every figure/diagram, write ONE description block exactly like this:
  [FIGURE: describe everything — shapes, dimensions, labeled axes, arrows, points,
   what physics concept it illustrates. Be detailed and specific.]
- Code: fenced blocks

Extract everything. Be faithful. No rephrasing. No skipping. OUTPUT = PAGE CONTENT ONLY.
"""

def extract_raw(client: Groq, base64_image: str, vision_model: str) -> str:
    response = client.chat.completions.create(
        model=vision_model,
        messages=[{"role": "user", "content": [
            {"type": "text",      "text": EXTRACTION_PROMPT},
            {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{base64_image}"}},
        ]}],
        temperature=0.1, max_completion_tokens=4096,
    )
    return response.choices[0].message.content


# ══════════════════════════════════════════════════════════════════════════════
#  STEP 2 — REWRITE PROMPT  (with SVG diagram generation)
# ══════════════════════════════════════════════════════════════════════════════

SVG_EXAMPLES = '''
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
SVG EXAMPLES — use these as templates for physics diagrams
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

EXAMPLE A — Vertical Circular Motion:
<svg xmlns="http://www.w3.org/2000/svg" width="320" height="320" viewBox="0 0 320 320">
  <style>
    text { font-family: "Georgia", serif; font-size: 13px; fill: #1a1a2e; }
    .label { font-size: 12px; fill: #374151; }
    .formula { font-size: 11px; fill: #6b46c1; font-style: italic; }
    .arrow { stroke: #1e40af; stroke-width: 2; fill: #1e40af; }
    .axis  { stroke: #9ca3af; stroke-width: 1; stroke-dasharray: 5,4; }
  </style>
  <!-- Circle path -->
  <circle cx="160" cy="155" r="100" fill="none" stroke="#3b82f6" stroke-width="2.5"/>
  <!-- Center -->
  <circle cx="160" cy="155" r="3" fill="#374151"/>
  <text x="165" y="152" class="label">O</text>
  <!-- Radius line to point P -->
  <line x1="160" y1="155" x2="230" y2="85" stroke="#6b7280" stroke-width="1.5" stroke-dasharray="4,3"/>
  <!-- Point P -->
  <circle cx="230" cy="85" r="5" fill="#ef4444"/>
  <text x="238" y="82" class="label" fill="#dc2626">P</text>
  <!-- Radius label -->
  <text x="188" y="108" class="formula">r</text>
  <!-- Weight arrow -->
  <line x1="230" y1="85" x2="230" y2="135" class="arrow" marker-end="url(#arr)"/>
  <text x="236" y="115" class="formula">mg</text>
  <!-- Velocity arrow (tangent) -->
  <line x1="230" y1="85" x2="175" y2="50" class="arrow" marker-end="url(#arr)" stroke="#10b981" fill="#10b981"/>
  <text x="162" y="42" class="formula" fill="#059669">v</text>
  <!-- Normal force arrow -->
  <line x1="230" y1="85" x2="270" y2="50" class="arrow" marker-end="url(#arr)" stroke="#f59e0b" fill="#f59e0b"/>
  <text x="272" y="46" class="formula" fill="#d97706">N</text>
  <!-- Bottom point label -->
  <circle cx="160" cy="255" r="4" fill="#374151"/>
  <text x="148" y="272" class="label">Bottom</text>
  <!-- Top point label -->
  <circle cx="160" cy="55" r="4" fill="#374151"/>
  <text x="148" y="45" class="label">Top</text>
  <!-- Arrow marker def -->
  <defs>
    <marker id="arr" markerWidth="8" markerHeight="8" refX="6" refY="3" orient="auto">
      <path d="M0,0 L0,6 L8,3 z" fill="inherit"/>
    </marker>
  </defs>
</svg>

EXAMPLE B — Thin Ring / Hollow Cylinder:
<svg xmlns="http://www.w3.org/2000/svg" width="280" height="220" viewBox="0 0 280 220">
  <style>
    text { font-family: "Georgia", serif; font-size: 12px; fill: #1a1a2e; }
    .formula { font-size: 12px; fill: #6b46c1; font-style: italic; }
    .dim { stroke: #9ca3af; stroke-width: 1; }
  </style>
  <!-- Outer ring -->
  <ellipse cx="140" cy="100" rx="90" ry="30" fill="none" stroke="#3b82f6" stroke-width="3"/>
  <!-- Axis line (vertical dashed) -->
  <line x1="140" y1="20" x2="140" y2="185" stroke="#9ca3af" stroke-width="1.5" stroke-dasharray="6,4"/>
  <text x="143" y="18" font-size="11" fill="#6b7280">axis</text>
  <!-- Radius arrow -->
  <line x1="140" y1="100" x2="228" y2="100" stroke="#ef4444" stroke-width="1.8" marker-end="url(#ra)"/>
  <text x="175" y="93" class="formula" fill="#dc2626">R</text>
  <!-- Mass label -->
  <text x="108" y="140" class="formula">M</text>
  <!-- Formula -->
  <text x="60" y="185" font-size="13" fill="#1e40af" font-style="italic">I = MR²</text>
  <!-- Arrow marker -->
  <defs>
    <marker id="ra" markerWidth="7" markerHeight="7" refX="5" refY="3.5" orient="auto">
      <path d="M0,0 L0,7 L7,3.5 z" fill="#ef4444"/>
    </marker>
  </defs>
</svg>

EXAMPLE C — Solid Disc / Cylinder:
<svg xmlns="http://www.w3.org/2000/svg" width="280" height="220" viewBox="0 0 280 220">
  <style>
    text { font-family: "Georgia", serif; font-size: 12px; fill: #1a1a2e; }
    .formula { font-size: 12px; fill: #6b46c1; font-style: italic; }
  </style>
  <!-- Filled disc -->
  <ellipse cx="140" cy="100" rx="90" ry="30" fill="#dbeafe" stroke="#3b82f6" stroke-width="2.5"/>
  <!-- Axis -->
  <line x1="140" y1="15" x2="140" y2="185" stroke="#9ca3af" stroke-width="1.5" stroke-dasharray="6,4"/>
  <text x="143" y="14" font-size="11" fill="#6b7280">axis</text>
  <!-- Radius -->
  <line x1="140" y1="100" x2="228" y2="100" stroke="#ef4444" stroke-width="1.8" marker-end="url(#rd)"/>
  <text x="178" y="93" class="formula" fill="#dc2626">R</text>
  <!-- Formula -->
  <text x="50" y="185" font-size="13" fill="#1e40af" font-style="italic">I = ½MR²</text>
  <defs>
    <marker id="rd" markerWidth="7" markerHeight="7" refX="5" refY="3.5" orient="auto">
      <path d="M0,0 L0,7 L7,3.5 z" fill="#ef4444"/>
    </marker>
  </defs>
</svg>

EXAMPLE D — Thin Rod (central axis):
<svg xmlns="http://www.w3.org/2000/svg" width="320" height="160" viewBox="0 0 320 160">
  <style>
    text { font-family: "Georgia", serif; font-size: 12px; fill: #1a1a2e; }
    .formula { font-size: 12px; fill: #6b46c1; font-style: italic; }
  </style>
  <!-- Rod -->
  <rect x="40" y="72" width="240" height="16" rx="3" fill="#bfdbfe" stroke="#3b82f6" stroke-width="2"/>
  <!-- Axis (vertical through center) -->
  <line x1="160" y1="15" x2="160" y2="145" stroke="#9ca3af" stroke-width="1.5" stroke-dasharray="6,4"/>
  <text x="163" y="14" font-size="11" fill="#6b7280">axis</text>
  <!-- Half-length arrows -->
  <line x1="160" y1="58" x2="280" y2="58" stroke="#ef4444" stroke-width="1.5" marker-end="url(#rl)" marker-start="url(#rls)"/>
  <text x="207" y="53" class="formula" fill="#dc2626">L/2</text>
  <line x1="160" y1="58" x2="40"  y2="58" stroke="#ef4444" stroke-width="1.5" marker-end="url(#rl)"/>
  <text x="85"  y="53" class="formula" fill="#dc2626">L/2</text>
  <!-- Formula -->
  <text x="75" y="148" font-size="13" fill="#1e40af" font-style="italic">I = ML²/12</text>
  <defs>
    <marker id="rl"  markerWidth="7" markerHeight="7" refX="5" refY="3.5" orient="auto">
      <path d="M0,0 L0,7 L7,3.5 z" fill="#ef4444"/>
    </marker>
    <marker id="rls" markerWidth="7" markerHeight="7" refX="2" refY="3.5" orient="auto-start-reverse">
      <path d="M0,0 L0,7 L7,3.5 z" fill="#ef4444"/>
    </marker>
  </defs>
</svg>

EXAMPLE E — Sphere (solid):
<svg xmlns="http://www.w3.org/2000/svg" width="260" height="260" viewBox="0 0 260 260">
  <style>
    text { font-family: "Georgia", serif; font-size: 12px; fill: #1a1a2e; }
    .formula { font-size: 12px; fill: #6b46c1; font-style: italic; }
  </style>
  <defs>
    <radialGradient id="sg" cx="40%" cy="35%">
      <stop offset="0%"   stop-color="#bfdbfe"/>
      <stop offset="100%" stop-color="#3b82f6"/>
    </radialGradient>
    <marker id="sa" markerWidth="7" markerHeight="7" refX="5" refY="3.5" orient="auto">
      <path d="M0,0 L0,7 L7,3.5 z" fill="#ef4444"/>
    </marker>
  </defs>
  <!-- Sphere -->
  <circle cx="130" cy="125" r="95" fill="url(#sg)" stroke="#2563eb" stroke-width="2"/>
  <!-- Equator ellipse -->
  <ellipse cx="130" cy="125" rx="95" ry="25" fill="none" stroke="#1d4ed8" stroke-width="1.2" stroke-dasharray="5,4"/>
  <!-- Axis -->
  <line x1="130" y1="15" x2="130" y2="235" stroke="#9ca3af" stroke-width="1.5" stroke-dasharray="5,4"/>
  <!-- Radius -->
  <line x1="130" y1="125" x2="213" y2="78" stroke="#ef4444" stroke-width="1.8" marker-end="url(#sa)"/>
  <text x="178" y="93" class="formula" fill="#dc2626">R</text>
  <!-- Formula -->
  <text x="50" y="248" font-size="13" fill="#1e40af" font-style="italic">I = ²⁄₅MR²</text>
</svg>

EXAMPLE F — Force / Free Body Diagram:
<svg xmlns="http://www.w3.org/2000/svg" width="240" height="260" viewBox="0 0 240 260">
  <style>
    text { font-family: "Georgia", serif; font-size: 12px; fill: #1a1a2e; }
    .formula { font-size: 12px; fill: #6b46c1; font-style: italic; }
  </style>
  <defs>
    <marker id="fa" markerWidth="8" markerHeight="8" refX="6" refY="4" orient="auto">
      <path d="M0,0 L0,8 L8,4 z" fill="inherit"/>
    </marker>
  </defs>
  <!-- Object box -->
  <rect x="90" y="105" width="60" height="50" rx="4" fill="#fef3c7" stroke="#d97706" stroke-width="2"/>
  <text x="108" y="135" font-size="11">body</text>
  <!-- Weight down -->
  <line x1="120" y1="155" x2="120" y2="215" stroke="#ef4444" stroke-width="2.2" marker-end="url(#fa)" style="fill:#ef4444"/>
  <text x="126" y="193" class="formula" fill="#dc2626">mg</text>
  <!-- Normal up -->
  <line x1="120" y1="105" x2="120" y2="45" stroke="#10b981" stroke-width="2.2" marker-end="url(#fa)" style="fill:#10b981"/>
  <text x="126" y="72" class="formula" fill="#059669">N</text>
  <!-- Friction left -->
  <line x1="90" y1="130" x2="30" y2="130" stroke="#8b5cf6" stroke-width="2.2" marker-end="url(#fa)" style="fill:#8b5cf6"/>
  <text x="32" y="122" class="formula" fill="#7c3aed">f</text>
  <!-- Applied right -->
  <line x1="150" y1="130" x2="210" y2="130" stroke="#f59e0b" stroke-width="2.2" marker-end="url(#fa)" style="fill:#f59e0b"/>
  <text x="173" y="122" class="formula" fill="#d97706">F</text>
</svg>
'''

def build_rewrite_prompt(raw_content: str, language: str) -> str:
    lang_name = SUPPORTED_LANGUAGES.get(language, language)

    if language == "en":
        lang_instr = (
            "Write in English. Completely rewrite all prose in your own words — "
            "same meaning, different phrasing and structure. Every concept must be covered."
        )
    else:
        lang_instr = (
            f"Write in {lang_name}. Translate AND rewrite all prose into fluent natural {lang_name}. "
            f"Every concept must be present. "
            f"Do NOT translate: LaTeX math, SVG code, Mermaid code, math symbols inside $ signs."
        )

    return f"""You are an expert academic writer, physics illustrator, and Markdown designer.

You have raw OCR-extracted content from a textbook page.
Rewrite it into a polished, visually rich Markdown document.

{lang_instr}

⛔ ABSOLUTE PROHIBITIONS — violating ANY of these is a critical failure:
- NO meta-commentary ("Here is the rewritten content...", "The content includes...", "Below is...")
- NO step-by-step process headers ("Step 1:", "Step 2:", "## 3:", "## 4:", etc.)
- NO "The final answer is:", "In conclusion:", "In summary:", or any closing remarks
- NO repeating the same content twice in any form
- NO restating what the raw content contains — just rewrite it directly
- NO listing/describing figures without generating actual SVG — every [FIGURE:] MUST become an SVG
- NO broken image links `![](...)`
- First character of your output MUST be actual document content (a heading or text)

━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
REWRITING RULES
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

1. CONTENT COVERAGE — non-negotiable:
   - Every concept, formula, term, and data point MUST appear EXACTLY ONCE
   - Tables must keep all rows and columns
   - All formulas must be preserved exactly
   - If the raw content already contains rewritten sections followed by raw extracted
     versions of the same content, OUTPUT ONLY the polished version — never duplicate.

2. LAYOUT REDESIGN — use your creativity:
   - Improve heading hierarchy (##, ###, ####)
   - Break dense paragraphs into clean bullet points where helpful
   - Add > 💡 **Key Insight:** callout blocks for important results
   - Add > ⚠️ **Note:** for important conditions or exceptions

3. MATH — preserve exactly:
   - Inline: $expression$
   - Block: $$ expression $$
   - Table cells: $...$ one line, \\vert not |

4. FIGURES — MOST IMPORTANT RULE:

   For every [FIGURE: ...] description in the raw content, generate a BEAUTIFUL inline SVG.
   If the raw content only describes figures in text ("A diagram shows...") with no [FIGURE:]
   tag, STILL generate an SVG based on the description.

   SVG DESIGN STANDARDS:
   - viewBox: use appropriate size (240–400 wide, 160–320 tall)
   - Colors: use this palette:
     • Shapes/outlines: #3b82f6 (blue) stroke, #dbeafe fill (light blue) or gradient
     • Force arrows: #ef4444 (red) for weight/gravity, #10b981 (green) for velocity,
                     #f59e0b (amber) for normal force, #8b5cf6 (purple) for friction
     • Axes/dashed: #9ca3af stroke, stroke-dasharray="5,4"
     • Dimension lines: #ef4444
     • Labels: font-family="Georgia,serif" font-size="12" fill="#1a1a2e"
     • Formulas in SVG: font-style="italic" fill="#6b46c1"
   - Always include: arrowhead markers via <defs><marker>
   - Label every important point, angle, dimension, force
   - Show rotation axis as dashed vertical/horizontal line
   - For 3D objects (spheres, cylinders): use ellipses for cross-sections
   - Add gradient fills for 3D effect where appropriate
   - Keep SVG clean and well-commented

   Wrap SVG in a blockquote with a caption:
   > **📐 Figure — [Title]**
   >
   > <svg ...> ... </svg>
   >
   > *[One-line caption explaining what the figure shows]*

{SVG_EXAMPLES}

   For FLOWCHARTS / PROCESSES — use Mermaid instead of SVG:
   ```mermaid
   graph TD
       A[Step 1] --> B[Step 2]
   ```

   For GRAPHS / PLOTS — use Mermaid xychart-beta if values are known, else SVG axes:
   ```mermaid
   xychart-beta
       title "Angular Velocity vs Time"
       x-axis [0, 1, 2, 3]
       y-axis 0 --> 30
       line [0, 10, 20, 30]
   ```

5. TABLES — full Markdown syntax, every row replicated

6. CODE — fenced blocks with language identifier

━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
STRICT MATH RULES
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
- Only $...$ and $$...$$ for math (never \\(...\\) or \\[...\\])

━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
RAW EXTRACTED CONTENT:
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

{raw_content}
"""

def rewrite_content(client: Groq, raw_content: str, language: str, rewrite_model: str) -> str:
    prompt = build_rewrite_prompt(raw_content, language)
    response = client.chat.completions.create(
        model=rewrite_model,
        messages=[{"role": "user", "content": prompt}],
        temperature=0.5,
        max_completion_tokens=8000,
    )
    return response.choices[0].message.content


# ══════════════════════════════════════════════════════════════════════════════
#  FILE BUILDERS
# ══════════════════════════════════════════════════════════════════════════════

def rebuild_chapter_file(chapter_dir: Path, lang_suffix: str) -> Path:
    pages_dir = chapter_dir / "pages"
    combined  = chapter_dir / f"combined{lang_suffix}.md"
    chunks = []
    if pages_dir.exists():
        for f in sorted(pages_dir.glob(f"page_*{lang_suffix}.md")):
            text = f.read_text(encoding="utf-8").strip()
            if text:
                chunks.append(text)
    combined.write_text("\n\n".join(chunks), encoding="utf-8")
    return combined

def rebuild_book_file(base_dir: Path, pdf_stem: str, lang_suffix: str) -> Path:
    book_file  = base_dir / f"{pdf_stem}_complete{lang_suffix}.md"
    all_chunks = []
    for ch_dir in sorted(d for d in base_dir.iterdir() if d.is_dir()):
        ch_combined = ch_dir / f"combined{lang_suffix}.md"
        if ch_combined.exists():
            text = ch_combined.read_text(encoding="utf-8").strip()
            if text:
                all_chunks.append(text)
    book_file.write_text("\n\n".join(all_chunks), encoding="utf-8")
    return book_file


# ══════════════════════════════════════════════════════════════════════════════
#  STATE
# ══════════════════════════════════════════════════════════════════════════════

def load_state(base_dir: Path) -> dict:
    f = base_dir / ".ocr_state.json"
    return json.loads(f.read_text()) if f.exists() else {"page_to_chapter": {}, "chapters": {}}

def save_state(base_dir: Path, state: dict):
    (base_dir / ".ocr_state.json").write_text(json.dumps(state, indent=2, ensure_ascii=False))


# --------------------------------------------------------------------------- #
#  MAIN PROCESSOR
# --------------------------------------------------------------------------- #

def process_pdf(
    pdf_path:          str,
    output_dir:        str   = None,
    language:          str   = "en",
    vision_model_key:  str   = DEFAULT_VISION_MODEL,
    rewrite_model_key: str   = DEFAULT_REWRITE_MODEL,
    delay:             float = 6.0,
    step_delay:        float = 3.0,
    detection_delay:   float = 2.0,
    retry_delay:       float = 30.0,
    max_retries:       int   = 3,
    resume:            bool  = True,
    skip_rewrite:      bool  = False,
    reset_detection:   bool  = False,
    force_single_chapter: bool = False,
):
    api_key = os.environ.get("GROQ_API_KEY","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")
    if not api_key:
        raise ValueError("GROQ_API_KEY environment variable not set.")

    if language not in SUPPORTED_LANGUAGES:
        for c, n in SUPPORTED_LANGUAGES.items(): print(f"  {c:6s}  {n}")
        raise ValueError(f"Unsupported language: {language}")

    vision_model  = VISION_MODELS.get(vision_model_key,  VISION_MODELS[DEFAULT_VISION_MODEL])
    rewrite_model = REWRITE_MODELS.get(rewrite_model_key, REWRITE_MODELS[DEFAULT_REWRITE_MODEL])
    lang_name     = SUPPORTED_LANGUAGES[language]
    lang_suffix   = f"_{language}" if language != "en" else ""

    client   = Groq(api_key=api_key)
    pdf_path = Path(pdf_path)
    if not pdf_path.exists():
        raise FileNotFoundError(f"PDF not found: {pdf_path}")

    base_dir = Path(output_dir) if output_dir else pdf_path.parent / pdf_path.stem
    base_dir.mkdir(parents=True, exist_ok=True)

    print(f"\n{'═'*64}")
    print(f"  📚  PDF OCR  —  2-Step Pipeline  |  SVG Diagrams")
    print(f"{'═'*64}")
    print(f"  📄 PDF      : {pdf_path.name}")
    print(f"  📁 Output   : {base_dir}")
    print(f"  🌐 Language : {lang_name} ({language})")
    print(f"  👁️  Step 1  : {vision_model.split('/')[-1]}  [extraction]")
    print(f"  ✍️  Step 2  : {'SKIPPED' if skip_rewrite else rewrite_model.split('/')[-1] + '  [rewrite + SVG diagrams]'}")
    print(f"  ⏱️  Delays  : {delay}s/page | {step_delay}s step1→2 | {detection_delay}s detect→extract")
    print(f"{'═'*64}\n")

    doc         = fitz.open(str(pdf_path))
    total_pages = len(doc)
    print(f"  📑 Total pages: {total_pages}\n")

    state           = load_state(base_dir)
    page_to_chapter = state["page_to_chapter"]
    chapters_meta   = state["chapters"]

    if reset_detection:
        # Wipe the chapter mapping so every page gets fresh detection.
        # Keep other state keys if any exist in future.
        page_to_chapter = {}
        chapters_meta   = {}
        save_state(base_dir, {"page_to_chapter": page_to_chapter, "chapters": chapters_meta})
        print(f"  🔄 Chapter detection state RESET — all pages will be re-detected.\n")

    current_chapter = page_to_chapter.get("_current", "chapter_00_introduction")

    for i, page in enumerate(doc):
        page_num = i + 1
        page_key = str(page_num)

        print(f"{'─'*64}")
        print(f"  Page {page_num}/{total_pages}")

        base64_img = pdf_page_to_base64(page)

        # ── Chapter detection ─────────────────────────────────────────────────
        if force_single_chapter:
            # Bypass detection entirely and use the output directory as the only "chapter"
            current_chapter = ""
            chapter_dir = base_dir 
        else:
            if page_key in page_to_chapter:
                current_chapter = page_to_chapter[page_key]
                print(f"  📂 Chapter : {current_chapter}  (from state)")
            else:
                print(f"  🔎 Detecting chapter...", end=" ", flush=True)
                det = detect_chapter(client, base64_img, vision_model)
                if det.get("is_new_chapter"):
                    folder_name = make_chapter_folder(
                        det.get("chapter_type"), det.get("chapter_number"),
                        det.get("chapter_title") or "untitled"
                    )
                    current_chapter = folder_name
                    chapters_meta[folder_name] = {
                        "title": det.get("chapter_title"), "number": det.get("chapter_number"),
                        "type":  det.get("chapter_type"),
                    }
                    print(f"📖 NEW → {folder_name}")
                else:
                    print(f"→ {current_chapter}")
                page_to_chapter[page_key]   = current_chapter
                page_to_chapter["_current"] = current_chapter
                save_state(base_dir, {"page_to_chapter": page_to_chapter, "chapters": chapters_meta})
                time.sleep(detection_delay)

            chapter_dir = base_dir / current_chapter

        pages_dir   = chapter_dir / "pages"
        pages_dir.mkdir(parents=True, exist_ok=True)

        raw_file   = pages_dir / f"page_{page_num:04d}_raw.md"
        final_file = pages_dir / f"page_{page_num:04d}{lang_suffix}.md"

        if resume and final_file.exists() and final_file.stat().st_size > 0:
            print(f"  ⏭️  Already complete, skipping.")
            continue

        # ── STEP 1: EXTRACTION ────────────────────────────────────────────────
        raw_content = None
        if resume and raw_file.exists() and raw_file.stat().st_size > 0:
            print(f"  ✅ Step 1 already saved, loading...")
            raw_content = raw_file.read_text(encoding="utf-8")
        else:
            print(f"  👁️  Step 1 — Extracting...", end=" ", flush=True)
            for attempt in range(1, max_retries + 1):
                try:
                    raw_content = extract_raw(client, base64_img, vision_model)
                    break
                except Exception as e:
                    err_str = str(e).lower()
                    wait = retry_delay * attempt if ("rate" in err_str or "429" in err_str) else 5.0
                    if attempt < max_retries:
                        print(f"\n    ⚠️  ({attempt}) {e}  — retrying in {wait}s...", end=" ", flush=True)
                        time.sleep(wait)
                    else:
                        raw_content = f"> ❌ Extraction failed: `{e}`\n"
            raw_file.write_text(raw_content, encoding="utf-8")
            print(f"saved ✅")

        # ── STEP 2: REWRITE + SVG ─────────────────────────────────────────────
        if skip_rewrite:
            final_content = post_process(raw_content)
            final_file.write_text(final_content, encoding="utf-8")
            print(f"  ⏩ Step 2 skipped.")
        else:
            time.sleep(step_delay)
            print(f"  ✍️  Step 2 — Rewriting + SVG diagrams...", end=" ", flush=True)
            rewritten = None
            for attempt in range(1, max_retries + 1):
                try:
                    rewritten = rewrite_content(client, raw_content, language, rewrite_model)
                    break
                except Exception as e:
                    err_str = str(e).lower()
                    wait = retry_delay * attempt if ("rate" in err_str or "429" in err_str) else 5.0
                    if attempt < max_retries:
                        print(f"\n    ⚠️  ({attempt}) {e}  — retrying in {wait}s...", end=" ", flush=True)
                        time.sleep(wait)
                    else:
                        print(f"\n    ❌ Rewrite failed, using raw output.")
                        rewritten = raw_content
            final_content = post_process(rewritten)
            final_file.write_text(final_content, encoding="utf-8")
            print(f"saved ✅")

        rebuild_chapter_file(chapter_dir, lang_suffix)
        rebuild_book_file(base_dir, pdf_path.stem, lang_suffix)

        if page_num < total_pages:
            print(f"  ⏸️  Waiting {delay}s ...\n")
            time.sleep(delay)

    doc.close()
    rebuild_book_file(base_dir, pdf_path.stem, lang_suffix)

    print(f"\n{'═'*64}")
    print(f"  ✅ Complete!")
    print(f"  🌐 Language : {lang_name} ({language})")
    print(f"  👁️  Step 1  : {vision_model.split('/')[-1]}")
    print(f"  ✍️  Step 2  : {rewrite_model.split('/')[-1]}")
    if chapters_meta:
        print(f"\n  📚 Chapters detected:")
        for folder, meta in sorted(chapters_meta.items()):
            print(f"     {meta.get('type','?')} {meta.get('number','?')}: {meta.get('title','—')}")
            print(f"       └─ {folder}/")
    print(f"\n  📁 Root  : {base_dir}")
    print(f"  📝 Book  : {base_dir}/{pdf_path.stem}_complete{lang_suffix}.md")
    print(f"{'═'*64}\n")
    return str(base_dir)


# ══════════════════════════════════════════════════════════════════════════════
#  CLI
# ══════════════════════════════════════════════════════════════════════════════

def main():
    parser = argparse.ArgumentParser(
        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
    )
    parser.add_argument("pdf", help="Path to the input PDF file")
    parser.add_argument("-o", "--output-dir", default=None,
        help="Output root folder (default: folder named after the PDF)")
    parser.add_argument("--vision-model", choices=list(VISION_MODELS.keys()),
        default=DEFAULT_VISION_MODEL,
        help=f"Step 1 vision model. Options: {list(VISION_MODELS.keys())}")
    parser.add_argument("--rewrite-model", choices=list(REWRITE_MODELS.keys()),
        default=DEFAULT_REWRITE_MODEL,
        help=f"Step 2 text model. Options: {list(REWRITE_MODELS.keys())}")
    parser.add_argument("--skip-rewrite", action="store_true",
        help="Run Step 1 extraction only, skip rewriting")
    parser.add_argument("-l", "--language", choices=list(SUPPORTED_LANGUAGES.keys()),
        default="en", metavar="LANG",
        help="Output language code (--list-languages to see all)")
    parser.add_argument("-d", "--delay",           type=float, default=6.0)
    parser.add_argument("--step-delay",            type=float, default=3.0,
        help="Delay between Step 1 and Step 2 on same page")
    parser.add_argument("--detection-delay",       type=float, default=2.0)
    parser.add_argument("--retry-delay",           type=float, default=30.0)
    parser.add_argument("--max-retries",           type=int,   default=3)
    parser.add_argument("--no-resume",             action="store_true")
    parser.add_argument("--reset-detection",       action="store_true",
        help="Wipe cached chapter detection and re-detect all pages (OCR/rewrite files are kept)")
    parser.add_argument("--list-languages",        action="store_true")
    parser.add_argument("--list-models",           action="store_true")

    args = parser.parse_args()

    if args.list_languages:
        print("\nSupported language codes:\n")
        for c, n in SUPPORTED_LANGUAGES.items(): print(f"  {c:6s}  {n}")
        print(); return

    if args.list_models:
        print("\nVision models (Step 1):\n")
        for k, v in VISION_MODELS.items():
            print(f"  {k:12s}  {v}{' ← default' if k == DEFAULT_VISION_MODEL else ''}")
        print("\nRewrite models (Step 2):\n")
        for k, v in REWRITE_MODELS.items():
            print(f"  {k:12s}  {v}{' ← default' if k == DEFAULT_REWRITE_MODEL else ''}")
        print(); return

    process_pdf(
        pdf_path          = args.pdf,
        output_dir        = args.output_dir,
        language          = args.language,
        vision_model_key  = args.vision_model,
        rewrite_model_key = args.rewrite_model,
        delay             = args.delay,
        step_delay        = args.step_delay,
        detection_delay   = args.detection_delay,
        retry_delay       = args.retry_delay,
        max_retries       = args.max_retries,
        resume            = not args.no_resume,
        skip_rewrite      = args.skip_rewrite,
        reset_detection   = args.reset_detection,
    )

if __name__ == "__main__":
    main()