# ============================================================
# SECTION: COMMON IMPORTS 
# ============================================================

import argparse
import json
import os
import re
import difflib
import string
import subprocess

from collections import Counter, defaultdict
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

import fitz  # PyMuPDF
import tiktoken
import openai
from dotenv import load_dotenv
from openai import OpenAI
from google.cloud import texttospeech

# ============================================================
# SECTION: chunking
# ============================================================

# ----------------------------
# Configuration & heuristics
# ----------------------------

NUMBERING_RE = re.compile(r"^\s*(\d+)(\.\d+){0,3}\s*[.:]?\s*")  # 3., 4.1., 4.1.1, etc.
PAGE_NUM_RE = re.compile(r"^\s*(Page\s+\d+|\d+)\s*$", re.IGNORECASE)

def is_all_caps(s: str) -> bool:
    letters = [c for c in s if c.isalpha()]
    if not letters:
        return False
    upper_ratio = sum(1 for c in letters if c.isupper()) / len(letters)
    return upper_ratio >= 0.8

def is_title_case_like(s: str) -> bool:
    """Loose check: proportion of words with initial uppercase."""
    tokens = [t for t in re.split(r"\s+", s) if t.isalpha()]
    if not tokens:
        return False
    upper_initial = sum(1 for t in tokens if t[0].isupper())
    return (upper_initial / len(tokens)) >= 0.6

def normalize_space(s: str) -> str:
    return re.sub(r"\s+", " ", s).strip()

def strip_visual_refs(s: str) -> str:
    """Lightly strip common visual references."""
    patterns = [
        r"\bsee\s+figure\b",
        r"\bas\s+shown\s+(above|below)\b",
        r"\bsee\s+table\b",
        r"\bdiagram\b",
        r"\bpage\s+\d+\b",
    ]
    if not s:
        return s
    return re.sub("|".join(patterns), "", s, flags=re.IGNORECASE)

def normalize_title(s: str) -> str:
    """Normalize title for robust matching across TOC and on-page headings."""
    if not isinstance(s, str):
        return ""
    s = s.lower().strip()
    s = s.translate(str.maketrans("", "", string.punctuation))
    s = re.sub(r"\s+", " ", s)
    return s

# ----------------------------
# Data models
# ----------------------------

@dataclass
class TextLine:
    text: str
    size_avg: float
    bold_any: bool
    font: str
    page: int
    block_idx: int
    line_idx: int
    y0: float
    y1: float

@dataclass
class Subsection:
    title: str
    start_page: int
    words: int = 0
    text: List[str] = field(default_factory=list)

@dataclass
class Section:
    title: str
    level: int     # 1=Chapter, 2=Section, 3=Subsection
    start_page: int
    end_page: Optional[int] = None
    words: int = 0                 # words from section preface/body (not subsections)
    text: List[str] = field(default_factory=list)  # preface/body (not subsections)
    subsections: List[Subsection] = field(default_factory=list)

@dataclass
class Chunk:
    id: str
    title: str
    source_titles: List[str]
    start_page: int
    end_page: int
    words: int
    minutes: float
    text: str


# ----------------------------
# TOC helpers (titles + pages)
# ----------------------------

def get_toc_entries(pdf_path: str) -> List[Tuple[int, str, int]]:
    """Returns a list of (level, normalized_title, page) for TOC/bookmarks (1-based pages)."""
    entries: List[Tuple[int, str, int]] = []
    doc = fitz.open(pdf_path)
    try:
        toc = doc.get_toc(simple=True)  # [(level, title, page), ...]
        for lvl, title, page in toc:
            if isinstance(title, str) and len(title.strip()) >= 3 and int(page) >= 1:
                entries.append((int(lvl), normalize_title(title), int(page)))
    except Exception:
        pass
    doc.close()
    return entries

def find_toc_anchor(
    toc_entries: List[Tuple[int, str, int]],
    query: Optional[str],
    fuzzy: bool = False,
    threshold: float = 0.80
) -> Optional[int]:
    """
    Return the 1-based page number for the TOC entry that matches `query`.
    - Exact match on normalized title first
    - Optional fuzzy match using difflib when exact fails
    """
    if not query:
        return None
    q_norm = normalize_title(query)

    for _lvl, title, page in toc_entries:
        if title == q_norm:
            return page

    if fuzzy:
        best_score, best_page = 0.0, None
        for _lvl, title, page in toc_entries:
            score = difflib.SequenceMatcher(None, q_norm, title).ratio()
            if score > best_score:
                best_score, best_page = score, page
        if best_page is not None and best_score >= threshold:
            return best_page

    return None

# ----------------------------
# PDF extraction (LINES + suppression)
# ----------------------------

def extract_lines(
    pdf_path: str,
    header_margin: float = 0.10,
    footer_margin: float = 0.10,
    repeat_threshold: float = 0.40
) -> List[TextLine]:
    """
    Extract text as aggregated LINES with style metadata.
    Also detect repeating phrases across pages for suppression.
    """
    doc = fitz.open(pdf_path)
    lines: List[TextLine] = []
    phrase_pages: defaultdict[str, set] = defaultdict(set)

    meta_title = normalize_space(doc.metadata.get("title") or "")

    for page_num, page in enumerate(doc):
        page_dict = page.get_text("dict")
        page_h = page.rect.height
        top_cut = page_h * header_margin
        bot_cut = page_h * (1.0 - footer_margin)

        for b_idx, block in enumerate(page_dict.get("blocks", [])):
            if "lines" not in block:
                continue

            for l_idx, line in enumerate(block["lines"]):
                line_text_parts: List[str] = []
                sizes: List[float] = []
                bold_flags: List[bool] = []
                font_name: str = ""

                y_top = None
                y_bot = None

                for span in line["spans"]:
                    txt = normalize_space(span.get("text", ""))
                    if not txt:
                        continue
                    line_text_parts.append(txt)
                    sizes.append(float(span.get("size", 0.0)))
                    bold_flags.append("Bold" in str(span.get("font", "")))
                    font_name = str(span.get("font", "")) or font_name

                    y0 = span.get("origin", [0, 0])[1]
                    y1 = y0 + span.get("size", 0)
                    y_top = y0 if y_top is None else min(y_top, y0)
                    y_bot = y1 if y_bot is None else max(y_bot, y1)

                if not line_text_parts:
                    continue

                text = normalize_space(" ".join(line_text_parts))
                size_avg = sum(sizes) / len(sizes) if sizes else 0.0
                bold_any = any(bold_flags)

                # Record phrase occurrences per page (for global suppression)
                if len(text) >= 6 and not PAGE_NUM_RE.match(text):
                    phrase_pages[text].add(page_num + 1)

                # Build line
                lines.append(
                    TextLine(
                        text=text,
                        size_avg=size_avg,
                        bold_any=bold_any,
                        font=font_name,
                        page=page_num + 1,  # 1-based index
                        block_idx=b_idx,
                        line_idx=l_idx,
                        y0=y_top or 0.0,
                        y1=y_bot or 0.0,
                    )
                )

    # Build suppression set:
    # - phrases that repeat on >= repeat_threshold of pages
    # - the PDF metadata Title (exact match)
    suppress_phrases = set()
    total_pages = len(doc)
    for phrase, pages_set in phrase_pages.items():
        if len(pages_set) / total_pages >= repeat_threshold:
            suppress_phrases.add(phrase)
    if meta_title:
        suppress_phrases.add(meta_title)

    doc.close()

    # Filter lines: drop suppressed and obvious page numbers
    filtered = [
        ln for ln in lines
        if ln.text not in suppress_phrases and not PAGE_NUM_RE.match(ln.text)
    ]
    return filtered

# ----------------------------
# Heading detection (TOC → numbering → fonts)
# ----------------------------

def estimate_body_size(lines: List[TextLine]) -> float:
    sizes = [round(ln.size_avg, 1) for ln in lines if len(ln.text) > 2]
    if not sizes:
        return 11.0
    size_counts = Counter(sizes)
    return size_counts.most_common(1)[0][0]

def collect_size_levels(lines: List[TextLine], body_size: float) -> Tuple[float, float]:
    bigger = [ln.size_avg for ln in lines if ln.size_avg > body_size and len(ln.text) >= 3]
    if not bigger:
        return (body_size + 3.0, body_size + 1.5)
    bigger_sorted = sorted(bigger)
    q90 = bigger_sorted[int(0.90 * (len(bigger_sorted) - 1))]
    q75 = bigger_sorted[int(0.75 * (len(bigger_sorted) - 1))]
    return (q90, q75)

def is_heading_like(text: str, ln: TextLine, body_size: float, in_toc: bool) -> bool:
    """
    Gate to prevent numbered list items from becoming headings.
    Must satisfy one of:
      - ln.size_avg >= body_size + 2.0
      - ln.bold_any
      - in_toc (exact TOC title match)
    And:
      - reasonable length (<= 20 words)
      - title-case-like OR all-caps
      - and avoid clear sentences (trailing ';' or '.' with many words)
    """
    gating_style = (ln.size_avg >= body_size + 2.0) or ln.bold_any or in_toc
    if not gating_style:
        return False

    words = re.split(r"\s+", text.strip())
    if len(words) == 0 or len(words) > 20:
        return False

    if re.search(r"[;.]$", text) and len(words) >= 5:
        return False

    return is_title_case_like(text) or is_all_caps(text)

def numbering_level(text: str, ln: TextLine, body_size: float, in_toc: bool, numbering_l1_is_chapter: bool) -> Optional[int]:
    """Map numbering depth to levels only if the line is heading-like."""
    if not NUMBERING_RE.match(text):
        return None
    if not is_heading_like(text, ln, body_size, in_toc):
        return None

    first_token = text.split()[0]
    dot_count = first_token.count(".")
    if numbering_l1_is_chapter:
        if dot_count == 1:
            return 1
        elif dot_count == 2:
            return 2
        else:
            return 3
    else:
        if dot_count == 1:
            return 2
        else:
            return 3

def heading_score(ln: TextLine, body_size: float, in_toc_titles: bool) -> float:
    score = 0.0
    score += max(0.0, (ln.size_avg - body_size)) / 2.0
    if ln.bold_any:
        score += 0.6
    if NUMBERING_RE.match(ln.text):
        score += 0.7
    if is_all_caps(ln.text):
        score += 0.3
    if in_toc_titles:
        score += 1.0
    return score

def classify_level_line(
    ln: TextLine,
    body_size: float,
    l1_threshold: float,
    l2_threshold: float,
    toc_map: Dict[str, int],
    numbering_l1_is_chapter: bool
) -> Optional[int]:
    """
    Determine Level for a line using priority:
      1) TOC exact title match → use TOC level (clamped to 1..3)
      2) Numbering depth (only if heading-like)
      3) Font-based heuristic thresholds
    """
    text_norm = normalize_title(ln.text)
    in_toc = text_norm in toc_map

    # 1) TOC level if exact title match
    if in_toc:
        lvl = max(1, min(3, toc_map[text_norm]))
        return lvl

    # 2) Numbering depth (guarded by heading-like)
    lvl_num = numbering_level(ln.text, ln, body_size, in_toc, numbering_l1_is_chapter)
    if lvl_num is not None:
        return lvl_num

    # 3) Font-based heuristic
    sc = heading_score(ln, body_size, in_toc)
    if ln.size_avg >= l1_threshold and sc >= 1.2 and is_heading_like(ln.text, ln, body_size, in_toc):
        return 1
    if ln.size_avg >= l2_threshold and sc >= 0.9 and is_heading_like(ln.text, ln, body_size, in_toc):
        return 2
    if (ln.size_avg >= body_size + 0.2 and sc >= 0.7 and is_heading_like(ln.text, ln, body_size, in_toc)):
        return 3

    return None

# ----------------------------
# Page cropping by anchors
# ----------------------------

def crop_lines_by_pages(
    lines: List[TextLine],
    start_page: Optional[int],
    end_page: Optional[int]
) -> List[TextLine]:
    """Keep only lines whose page falls within [start_page, end_page] bounds."""
    def in_bounds(p: int) -> bool:
        if start_page is not None and p < start_page:
            return False
        if end_page is not None and p > end_page:
            return False
        return True

    return [ln for ln in lines if in_bounds(ln.page)]

# ----------------------------
# Sections & subsections (from lines)
# ----------------------------

def build_sections_from_lines(
    lines: List[TextLine],
    body_size: float,
    l1: float,
    l2: float,
    toc_map: Dict[str, int],
    numbering_l1_is_chapter: bool
) -> List[Section]:
    sections: List[Section] = []
    current: Optional[Section] = None
    current_sub: Optional[Subsection] = None

    for ln in lines:
        level = classify_level_line(ln, body_size, l1, l2, toc_map, numbering_l1_is_chapter)

        if level in (1, 2):
            # Close previous
            if current:
                if current_sub:
                    current.subsections.append(current_sub)
                    current_sub = None
                current.end_page = current.end_page or current.start_page
                sections.append(current)

            # New section
            current = Section(
                title=normalize_space(ln.text),
                level=level,
                start_page=ln.page,
                end_page=None,
                words=0,
                text=[],
                subsections=[]
            )
            current_sub = None

        elif level == 3 and current:
            # Start new subsection
            if current_sub:
                current.subsections.append(current_sub)
            current_sub = Subsection(
                title=normalize_space(ln.text),
                start_page=ln.page,
                words=0,
                text=[]
            )

        else:
            # Body line
            if current is None:
                current = Section(
                    title="(Intro)",
                    level=2,
                    start_page=ln.page,
                    end_page=None,
                    words=0,
                    text=[],
                    subsections=[]
                )

            clean = strip_visual_refs(ln.text)
            if clean:
                if current_sub is not None:
                    current_sub.text.append(clean)
                    current_sub.words += len(clean.split())
                else:
                    current.text.append(clean)
                    current.words += len(clean.split())
                current.end_page = ln.page

    # Close trailing
    if current:
        if current_sub:
            current.subsections.append(current_sub)
        current.end_page = current.end_page or current.start_page
        sections.append(current)

    # Merge accidental “split headings”
    merged: List[Section] = []
    skip_next = False
    for i in range(len(sections)):
        if skip_next:
            skip_next = False
            continue
        s = sections[i]
        if i + 1 < len(sections):
            n = sections[i + 1]
            if s.words == 0 and n.level == s.level and s.start_page == n.start_page and len(n.text) == 0:
                s.title = normalize_space(f"{s.title} {n.title}")
                s.end_page = max(s.end_page or s.start_page, n.end_page or n.start_page)
                skip_next = True
        merged.append(s)

    return merged

# ----------------------------
# Helpers: section totals (text + words)
# ----------------------------

def section_total_words(sec: Section) -> int:
    """Words in section preface/body + all subsections."""
    return sec.words + sum(ss.words for ss in sec.subsections)

def section_full_text(sec: Section) -> str:
    """Concatenate section body and all subsection texts in order."""
    parts: List[str] = []
    if sec.text:
        parts.append(" ".join(sec.text))
    for ss in sec.subsections:
        if ss.text:
            parts.append(" ".join(ss.text))
    return normalize_space(" ".join(parts))

def minutes_from_section(sec: Section, wpm: int) -> float:
    return round(section_total_words(sec) / max(120, wpm), 2)

# ----------------------------
# Grouping by boundary (TOC or detected)
# ----------------------------

def build_groups_from_toc_pages(
    sections: List[Section],
    toc_entries: List[Tuple[int, str, int]],
    toc_level: int,
    start_page: Optional[int],
    end_page: Optional[int]
) -> List[List[Section]]:
    """
    Build groups using TOC page anchors at the specified level.
    Each group contains sections whose start_page falls between consecutive TOC pages.
    """
    pages = sorted({page for lvl, _title, page in toc_entries if lvl == toc_level})
    if start_page:
        pages = [p for p in pages if p >= start_page]
    if end_page:
        pages = [p for p in pages if p <= end_page]
    if not pages:
        return [sections] if sections else []

    last_page = end_page or max(s.end_page or s.start_page for s in sections)
    groups: List[List[Section]] = []
    for i, p_start in enumerate(pages):
        p_end = (pages[i + 1] - 1) if (i + 1 < len(pages)) else last_page
        group = [s for s in sections if p_start <= s.start_page <= p_end]
        if group:
            groups.append(group)
    return groups

def group_by_level_detected(sections: List[Section], boundary_level: int) -> List[List[Section]]:
    groups: List[List[Section]] = []
    current: List[Section] = []

    for sec in sections:
        if sec.level == boundary_level:
            if current:
                groups.append(current)
            current = [sec]
        else:
            if not current:
                current = [sec]
            else:
                current.append(sec)

    if current:
        groups.append(current)

    return groups

# ----------------------------
# Audio length estimate
# ----------------------------

def minutes_from_text(words: int, wpm: int = 155) -> float:
    return round(words / max(120, wpm), 2)

# ----------------------------
# Chunking
# ----------------------------

def split_by_subsections(section: Section, target: int, max_mins: int, wpm: int) -> List[Chunk]:
    """
    Split a long section into sub-chunks using L3 subsections as boundaries; greedy towards target minutes.
    """
    chunks: List[Chunk] = []

    # Build subsection payloads including any body text before the first subsection
    units: List[Tuple[str, int, int, str]] = []  # (title, words, page, text)
    preface_text = normalize_space(" ".join(section.text))
    preface_words = len(preface_text.split())
    if preface_words:
        units.append((f"{section.title} – Overview", preface_words, section.start_page, preface_text))

    for ss in section.subsections:
        ss_text = normalize_space(" ".join(ss.text))
        units.append((f"{section.title} – {ss.title}", len(ss_text.split()), ss.start_page, ss_text))

    buffer_titles: List[str] = []
    buffer_texts: List[str] = []
    buffer_words: int = 0
    buffer_start_page: int = section.start_page

    def flush_chunk(end_page: int):
        nonlocal buffer_titles, buffer_texts, buffer_words, buffer_start_page
        if buffer_words == 0:
            return
        minutes = minutes_from_text(buffer_words, wpm)
        chunk = Chunk(
            id="",
            title=" | ".join(buffer_titles),
            source_titles=buffer_titles.copy(),
            start_page=buffer_start_page,
            end_page=end_page,
            words=buffer_words,
            minutes=minutes,
            text=normalize_space(" ".join(buffer_texts)),
        )
        chunks.append(chunk)
        buffer_titles, buffer_texts = [], []
        buffer_words = 0
        buffer_start_page = end_page

    for title, words, start_page, text in units:
        unit_minutes = minutes_from_text(words, wpm)
        if unit_minutes > max_mins:
            sentences = re.split(r"(?<=[.!?])\s+", text)
            sent_buf: List[str] = []
            sent_words = 0
            for s in sentences:
                sent_buf.append(s)
                sent_words += len(s.split())
                if minutes_from_text(sent_words, wpm) >= target:
                    buffer_titles.append(title)
                    buffer_texts.append(normalize_space(" ".join(sent_buf)))
                    buffer_words += sent_words
                    flush_chunk(start_page)
                    sent_buf, sent_words = [], 0
            if sent_words:
                buffer_titles.append(title)
                buffer_texts.append(normalize_space(" ".join(sent_buf)))
                buffer_words += sent_words
        else:
            buffer_titles.append(title)
            buffer_texts.append(text)
            buffer_words += words
            if minutes_from_text(buffer_words, wpm) >= target:
                flush_chunk(start_page)

    flush_chunk(section.end_page or section.start_page)
    return chunks

def chunk_sections(sections: List[Section], min_mins: int, target_mins: int, max_mins: int, wpm: int) -> List[Chunk]:
    """
    Chunk across sections with audio-aware bounds, using total words (section + subsections).
    """
    chunks: List[Chunk] = []
    buffer: Optional[Chunk] = None

    def new_chunk_from_section(sec: Section) -> Chunk:
        text = section_full_text(sec)
        words = section_total_words(sec)
        minutes = minutes_from_text(words, wpm)
        title = sec.title
        return Chunk(
            id="",
            title=title,
            source_titles=[sec.title],
            start_page=sec.start_page,
            end_page=sec.end_page or sec.start_page,
            words=words,
            minutes=minutes,
            text=text,
        )

    for sec in sections:
        sec_words_total = section_total_words(sec)
        sec_minutes_total = minutes_from_text(sec_words_total, wpm)

        # Split by subsections when the total section length exceeds max
        if sec_minutes_total > max_mins and sec.subsections:
            sub_chunks = split_by_subsections(sec, target_mins, max_mins, wpm)
            for ch in sub_chunks:
                if ch.minutes < min_mins and buffer:
                    merged_minutes = buffer.minutes + ch.minutes
                    if merged_minutes <= max_mins:
                        buffer.text = normalize_space(buffer.text + " " + ch.text)
                        buffer.words += ch.words
                        buffer.minutes = round(buffer.minutes + ch.minutes, 2)
                        buffer.source_titles += ch.source_titles
                        buffer.end_page = ch.end_page
                        buffer.title = buffer.title + " | " + ch.title
                        chunks.append(buffer)
                        buffer = None
                    else:
                        chunks.append(buffer)
                        buffer = ch
                elif ch.minutes < min_mins and not buffer:
                    buffer = ch
                else:
                    if buffer:
                        chunks.append(buffer)
                        buffer = None
                    chunks.append(ch)
            continue

        ch = new_chunk_from_section(sec)

        if ch.minutes < min_mins:
            if buffer:
                merged_minutes = buffer.minutes + ch.minutes
                if merged_minutes <= max_mins:
                    buffer.text = normalize_space(buffer.text + " " + ch.text)
                    buffer.words += ch.words
                    buffer.minutes = round(buffer.minutes + ch.minutes, 2)
                    buffer.source_titles += ch.source_titles
                    buffer.end_page = ch.end_page
                    buffer.title = buffer.title + " | " + ch.title
                else:
                    chunks.append(buffer)
                    buffer = ch
            else:
                buffer = ch
        elif ch.minutes > max_mins and sec.subsections:
            sub_chunks = split_by_subsections(sec, target_mins, max_mins, wpm)
            if buffer:
                chunks.append(buffer)
                buffer = None
            chunks.extend(sub_chunks)
        else:
            if buffer:
                if buffer.minutes + ch.minutes <= max_mins:
                    buffer.text = normalize_space(buffer.text + " " + ch.text)
                    buffer.words += ch.words
                    buffer.minutes = round(buffer.minutes + ch.minutes, 2)
                    buffer.source_titles += ch.source_titles
                    buffer.end_page = ch.end_page
                    buffer.title = buffer.title + " | " + ch.title
                    chunks.append(buffer)
                    buffer = None
                else:
                    chunks.append(buffer)
                    buffer = None
                    chunks.append(ch)
            else:
                chunks.append(ch)

    if buffer:
        chunks.append(buffer)

    for i, c in enumerate(chunks, start=1):
        c.id = f"chunk_{i:03d}"

    return chunks

def chunk_sections_with_groups(groups: List[List[Section]], min_mins: int, target_mins: int, max_mins: int, wpm: int, verbose: bool) -> List[Chunk]:
    chunks: List[Chunk] = []
    if verbose:
        print(f"• Hard boundary using {len(groups)} groups")
        for i, g in enumerate(groups, start=1):
            start_page = g[0].start_page
            end_page = g[-1].end_page or g[-1].start_page
            g_minutes = sum(minutes_from_section(s, wpm) for s in g)
            print(f"  Group {i}: pages {start_page}-{end_page}, ~{round(g_minutes, 1)} min")

    for g in groups:
        group_chunks = chunk_sections(g, min_mins, target_mins, max_mins, wpm)
        chunks.extend(group_chunks)

    for i, c in enumerate(chunks, start=1):
        c.id = f"chunk_{i:03d}"

    return chunks

# ----------------------------
# I/O helpers
# ----------------------------

def save_chunks(chunks: List[Chunk], outdir: str, pdf_basename: str, wpm: int) -> None:
    os.makedirs(outdir, exist_ok=True)
    index: List[Dict[str, Any]] = []
    for ch in chunks:
        safe_title = re.sub(r'[^A-Za-z0-9]+', '_', ch.title).strip('_')
        fname = f"{ch.id}__{safe_title[:80]}.txt"
        path = os.path.join(outdir, fname)
        with open(path, "w", encoding="utf-8") as f:
            f.write(ch.text)

        index.append({
            "id": ch.id,
            "title": ch.title,
            "source_titles": ch.source_titles,
            "file": fname,
            "pdf": pdf_basename,
            "start_page": ch.start_page,
            "end_page": ch.end_page,
            "word_count": ch.words,
            "minutes_estimate": ch.minutes,
            "wpm": wpm,
        })

    with open(os.path.join(outdir, "index.json"), "w", encoding="utf-8") as f:
        json.dump(index, f, indent=2, ensure_ascii=False)

def save_sections_debug(sections: List[Section], outdir: str) -> None:
    os.makedirs(outdir, exist_ok=True)
    debug = []
    for s in sections:
        debug.append({
            "title": s.title,
            "level": s.level,
            "start_page": s.start_page,
            "end_page": s.end_page,
            "words": s.words,
            "subsections": [
                {"title": ss.title, "start_page": ss.start_page, "words": ss.words}
                for ss in s.subsections
            ],
        })
    with open(os.path.join(outdir, "sections.debug.json"), "w", encoding="utf-8") as f:
        json.dump(debug, f, indent=2, ensure_ascii=False)




# ===========================
# chunking
# ===========================

def run_chunk(argv: Optional[List[str]] = None) -> None:
    ap = argparse.ArgumentParser(description="Audio-friendly Level 2 chunking (PyMuPDF)")
    ap.add_argument("pdf", help="Path to the well-structured course notes PDF")
    ap.add_argument("--outdir", default="./input_chunks", help="Output directory")
    ap.add_argument("--min", type=int, default=6, help="Minimum minutes per chunk")
    ap.add_argument("--target", type=int, default=16, help="Target minutes per chunk")
    ap.add_argument("--max", type=int, default=25, help="Maximum minutes per chunk")
    ap.add_argument("--wpm", type=int, default=155, help="Words per minute for estimate")

    ap.add_argument("--repeat-threshold", type=float, default=0.40,
                    help="Global suppression threshold for repeated phrases (fraction of pages)")
    ap.add_argument("--header-margin", type=float, default=0.10,
                    help="Top margin ratio for header detection (for stats)")
    ap.add_argument("--footer-margin", type=float, default=0.10,
                    help="Bottom margin ratio for footer detection (for stats)")

    ap.add_argument("--debug-sections", action="store_true", help="Write sections.debug.json")
    ap.add_argument("--verbose", action="store_true", help="Verbose diagnostics")

    ap.add_argument("--toc-start", type=str, default=None,
                    help="Ignore everything before this TOC entry (exact or fuzzy match)")
    ap.add_argument("--toc-end", type=str, default=None,
                    help="Ignore everything after this TOC entry (exact or fuzzy match)")
    ap.add_argument("--toc-fuzzy", action="store_true", help="Allow fuzzy matching for --toc-start/--toc-end")
    ap.add_argument("--toc-threshold", type=float, default=0.80, help="Fuzzy match threshold")

    ap.add_argument("--top-boundary-hard", action="store_true",
                    help="Force chunk breaks at the end of the chosen boundary level.")
    ap.add_argument("--boundary-level", type=int, choices=[1, 2, 3],
                    help="Heading level to enforce as hard boundary (1=Chapter, 2=Section, 3=Subsection).")
    ap.add_argument("--boundary-source", type=str, choices=["toc", "detected"], default="toc",
                    help="Use TOC pages or detected sections to build boundary groups.")

    ap.add_argument("--numbering-l1-is-chapter", action="store_true",
                    help="Treat single-level numbering (e.g., '3.') as Level-1 (chapter). Otherwise it's Level-2.")

    args = ap.parse_args(argv)

    pdf_path = args.pdf
    pdf_basename = os.path.basename(pdf_path)

    # --- TOC ---
    toc_entries = get_toc_entries(pdf_path)
    if args.verbose:
        if toc_entries:
            print(f"• TOC/bookmarks detected: {len(toc_entries)} entries.")
            preview = "; ".join([t for _lvl, t, _page in toc_entries[:8]])
            if preview:
                print(f"  TOC sample (normalized): {preview}")
        else:
            print("• No TOC/bookmarks detected.")

    toc_map: Dict[str, int] = {title: lvl for (lvl, title, _page) in toc_entries}

    start_page_anchor = find_toc_anchor(toc_entries, args.toc_start, args.toc_fuzzy, args.toc_threshold)
    end_page_anchor   = find_toc_anchor(toc_entries, args.toc_end,   args.toc_fuzzy, args.toc_threshold)
    if args.verbose:
        print(f"• TOC start anchor: {args.toc_start!r} -> page {start_page_anchor}")
        print(f"• TOC end   anchor: {args.toc_end!r}   -> page {end_page_anchor}")

    # --- Extract LINES + suppression ---
    lines = extract_lines(pdf_path,
                          header_margin=args.header_margin,
                          footer_margin=args.footer_margin,
                          repeat_threshold=args.repeat_threshold)

    if start_page_anchor or end_page_anchor:
        if start_page_anchor and end_page_anchor and start_page_anchor > end_page_anchor:
            if args.verbose:
                print(f"• Warning: start anchor page ({start_page_anchor}) > end anchor page ({end_page_anchor}). Swapping.")
            start_page_anchor, end_page_anchor = end_page_anchor, start_page_anchor
        lines = crop_lines_by_pages(lines, start_page_anchor, end_page_anchor)
        if args.verbose:
            if lines:
                min_page = min(ln.page for ln in lines)
                max_page = max(ln.page for ln in lines)
                print(f"• Cropped lines to pages [{min_page}, {max_page}] using TOC anchors.")
            else:
                print("• Cropping by TOC anchors removed all lines (check anchors).")

    # --- Heading thresholds & sections ---
    body_size = estimate_body_size(lines)
    l1, l2 = collect_size_levels(lines, body_size)

    if args.verbose:
        print(f"• Body font size (mode): {body_size}")
        print(f"• Heading thresholds -> L1: {l1:.2f}, L2: {l2:.2f}")
        print(f"• Total lines: {len(lines)}")

    sections = build_sections_from_lines(
        lines, body_size, l1, l2, toc_map, args.numbering_l1_is_chapter
    )

    if args.debug_sections:
        save_sections_debug(sections, args.outdir)

    if args.verbose:
        levels = Counter(s.level for s in sections)
        print(f"• Section level distribution: {dict(levels)}")
        sample_l1 = [s.title for s in sections if s.level == 1][:5]
        sample_l2 = [s.title for s in sections if s.level == 2][:5]
        sample_l3 = [s.title for s in sections if s.level == 3][:5]
        if sample_l1: print("  L1 samples:", "; ".join(sample_l1))
        if sample_l2: print("  L2 samples:", "; ".join(sample_l2))
        if sample_l3: print("  L3 samples:", "; ".join(sample_l3))

        sec_minutes = sum(minutes_from_section(s, args.wpm) for s in sections)
        print(f"• Aggregate section duration estimate: {round(sec_minutes, 1)} minutes (pre-chunking)")

    # --- Build boundary groups (TOC pages or detected) ---
    if args.top_boundary_hard:
        if args.boundary_source == "toc" and toc_entries:
            groups = build_groups_from_toc_pages(
                sections,
                toc_entries,
                toc_level=args.boundary_level if args.boundary_level else (1 if any(lvl == 1 for lvl, _, _ in toc_entries) else 2),
                start_page=start_page_anchor,
                end_page=end_page_anchor
            )
            chunks = chunk_sections_with_groups(groups, args.min, args.target, args.max, args.wpm, verbose=args.verbose)
        else:
            boundary_level = args.boundary_level if args.boundary_level else (1 if any(s.level == 1 for s in sections) else 2)
            groups = group_by_level_detected(sections, boundary_level)
            chunks = chunk_sections_with_groups(groups, args.min, args.target, args.max, args.wpm, verbose=args.verbose)
    else:
        chunks = chunk_sections(sections, args.min, args.target, args.max, args.wpm)

    # --- Save ---
    save_chunks(chunks, args.outdir, pdf_basename, args.wpm)

    print(f"✓ Created {len(chunks)} chunks in {args.outdir}")
    total_minutes = round(sum(c.minutes for c in chunks), 1)
    print(f"  Total duration (est.): {total_minutes} minutes at {args.wpm} WPM")
    print("  Next step: LLM rewrite -> TTS -> MP3")

    if args.verbose:
        print("• Chunk summaries (first 10):")
        for ch in chunks[:10]:
            print(f"  {ch.id}: {ch.title} | {ch.minutes} min | pages {ch.start_page}-{ch.end_page}")
        if len(chunks) > 10:
            print(f"  ...(and {len(chunks) - 10} more)")


# ============================================================
# SECTION: generate
# ============================================================

load_dotenv()

openai_client = OpenAI(
    api_key=os.getenv("OPENAI_API_KEY")
)

if not openai_client.api_key:
    raise RuntimeError("OPENAI_API_KEY not found. Check your .env file.")

# -------------------------------------------------
# Paths
# -------------------------------------------------
INPUT_DIR_GENERATE = Path("input_chunks")
OUTPUT_DIR_GENERATE = Path("output_scripts")
OUTPUT_DIR_GENERATE.mkdir(parents=True, exist_ok=True)

# -------------------------------------------------
# Full GPT-5.2 Gold-Standard Prompt
# -------------------------------------------------

SYSTEM_PROMPT = """
You are a senior instructional designer and educator specialising in audio-first learning for constrained-attention environments, such as listening while driving.

Your sole objective is to transform formal course notes into a clear, cognitively efficient spoken explanation that can be understood and retained without any visual reference.

You are not writing prose, lecture notes, or summaries.
You are producing a script that will be read aloud verbatim by a text-to-speech system.

Prioritise:
1. conceptual clarity
2. structure for listening
3. accuracy
4. retention

De-prioritise stylistic flourish, rhetorical creativity, or elegance.

Assume the listener:
- is driving and cannot pause frequently
- cannot see figures, tables, equations, or formatting
- is revising for assessment
- has limited working memory
- benefits from explicit signposting and pacing

Assume this script is one chunk in a larger audio curriculum and should feel consistent in tone and structure with other chunks.

Input may include:
- headings and subheadings
- dense academic language
- references to figures, tables, or equations
- material not suitable for direct audio delivery

Task Instructions:

1. Audio-first reformulation
   - Rewrite content into natural spoken language
   - Prefer short, direct sentences
   - Explain ideas sequentially, one concept at a time
   - Make relationships explicit (cause, effect, contrast, dependency)

2. Self-contained delivery
   - Remove or rewrite references to visuals (figures, diagrams, tables)
   - Remove page numbers, footnotes, "as shown above/below"
   - If a visual element is conceptually important, describe it verbally

3. Spoken structure (mandatory)
   - Opening orientation: clearly state what this section is about
   - Core explanation: introduce and explain key concepts, definitions, relationships
   - Why this matters: explain implications, applications, relevance
   - Common misunderstandings/pitfalls: only if it clarifies
   - Brief wrap-up: restate the central takeaway succinctly

4. Cognitive load management
   - Introduce no more than one new idea at a time
   - Avoid long lists unless each item is explained
   - Avoid nested clauses or compound sentences
   - Prefer clarity over completeness

5. Depth calibration
   - Do not oversimplify core concepts
   - Do not attempt to preserve every detail
   - Preserve the conceptual spine of the material

6. Audio suitability constraints
   - No bullet points
   - No headings or formatting
   - No equations unless clearly explained verbally
   - No meta commentary or source references

7. Output Requirements
   - Output ONLY the final audio script
   - Do not include explanations of changes, summaries, or commentary
   - Assume output will go directly to a TTS engine

8. Length guidance
   - Target spoken length: 10-25 minutes
   - If input is too long, prioritise definitions, relationships, causality
   - Do not pad content unnecessarily
"""

USER_PROMPT_TEMPLATE = """
BEGIN COURSE TEXT:

{course_text}
"""

# -------------------------------------------------
# Model configuration
# -------------------------------------------------
MODEL_NAME = "gpt-5.2"
TEMPERATURE = 0.3

# GPT-5.2 pricing per million tokens
COST_PER_MILLION_INPUT = 1.75
COST_PER_MILLION_OUTPUT = 14.0

# -------------------------------------------------
# Token encoder
# -------------------------------------------------
ENCODER = tiktoken.get_encoding("cl100k_base")

def count_tokens(text: str) -> int:
    return len(ENCODER.encode(text))

def estimate_cost(input_tokens: int, output_tokens: int) -> float:
    """
    Returns estimated cost in USD
    """
    cost = (input_tokens / 1_000_000) * COST_PER_MILLION_INPUT
    cost += (output_tokens / 1_000_000) * COST_PER_MILLION_OUTPUT
    return cost

def generate_audio_script(input_text: str) -> (str, int, int):
    response = openai_client.chat.completions.create(
        model=MODEL_NAME,
        temperature=TEMPERATURE,
        messages=[
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": USER_PROMPT_TEMPLATE.format(course_text=input_text)}
        ],
    )
    output_text = response.choices[0].message.content.strip()

    # count tokens for cost estimation
    input_tokens = count_tokens(input_text)
    output_tokens = count_tokens(output_text)
    return output_text, input_tokens, output_tokens


def run_generate():
    txt_files = sorted(INPUT_DIR_GENERATE.glob("*.txt"))
    if not txt_files:
        raise RuntimeError("No .txt files found in input_chunks directory.")

    total_input_tokens = 0
    total_output_tokens = 0
    total_estimated_cost = 0.0

    for txt_file in txt_files:
        print(f"Processing: {txt_file.name}")

        with txt_file.open("r", encoding="utf-8") as f:
            course_text = f.read().strip()

        if not course_text:
            print(f"Skipping empty file: {txt_file.name}")
            continue

        try:
            audio_script, input_tokens, output_tokens = generate_audio_script(course_text)

            # update totals
            total_input_tokens += input_tokens
            total_output_tokens += output_tokens
            chunk_cost = estimate_cost(input_tokens, output_tokens)
            total_estimated_cost += chunk_cost

            # save output
            output_path = OUTPUT_DIR_GENERATE / txt_file.name
            with output_path.open("w", encoding="utf-8") as out:
                out.write(audio_script)

            print(f"Saved output to: {output_path.name}")
            print(f"Chunk tokens: input={input_tokens}, output={output_tokens}")
            print(f"Chunk estimated cost: ${chunk_cost:.4f}")
            print(f"Running total estimated cost: ${total_estimated_cost:.4f}\n")

        except Exception as e:
            print(f"Error processing {txt_file.name}: {e}")

    print("--------------------------------------------------")
    print(f"All chunks processed: {len(txt_files)} files")
    print(f"Total input tokens: {total_input_tokens}")
    print(f"Total output tokens: {total_output_tokens}")
    print(f"Total estimated cost: ${total_estimated_cost:.4f}")
    print("--------------------------------------------------")



# ============================================================
# SECTION: tts 
# ============================================================

# -------------------------------------------------
# Paths
# -------------------------------------------------
INPUT_DIR = Path("output_scripts")  # GPT audio scripts
OUTPUT_DIR = Path("audio_output")
OUTPUT_DIR.mkdir(parents=True, exist_ok=True)

# -------------------------------------------------
# Google Cloud TTS client
# -------------------------------------------------
tts_client = texttospeech.TextToSpeechClient()

# -------------------------------------------------
# Voice & audio configuration
# -------------------------------------------------
VOICE_PARAMS = texttospeech.VoiceSelectionParams(
    language_code="en-AU",
    name="en-AU-Chirp3-HD-Laomedeia",
    ssml_gender=texttospeech.SsmlVoiceGender.NEUTRAL
)

AUDIO_CONFIG = texttospeech.AudioConfig(
    audio_encoding=texttospeech.AudioEncoding.MP3,
    speaking_rate=1.0,
    pitch=0.0
)

# -------------------------------------------------
# Helper functions
# -------------------------------------------------
def split_text_into_chunks(text, max_bytes=5000):
    """
    Split text into chunks <= max_bytes in UTF-8.
    Returns a list of strings.
    """
    chunks = []
    start = 0
    while start < len(text):
        end = start + max_bytes
        chunk = text[start:end]
        # Back off if UTF-8 byte length exceeds max_bytes
        while len(chunk.encode('utf-8')) > max_bytes:
            end -= 100
            chunk = text[start:end]
        chunks.append(chunk)
        start = end
    return chunks

def synthesize_text_chunk(text_chunk, output_path: Path, chunk_number: int, total_chunks: int):
    """
    Convert a single text chunk to MP3 and save it.
    Prints verbose updates.
    """
    print(f"  Synthesizing chunk {chunk_number}/{total_chunks} -> {output_path.name}")
    synthesis_input = texttospeech.SynthesisInput(text=text_chunk)
    response = tts_client.synthesize_speech(
        input=synthesis_input,
        voice=VOICE_PARAMS,
        audio_config=AUDIO_CONFIG
    )
    with open(output_path, "wb") as out_file:
        out_file.write(response.audio_content)
    print(f"    Saved: {output_path.name}")


def run_tts():
    txt_files = sorted(INPUT_DIR.glob("*.txt")) 

    if not txt_files:
        print("No text files found in output_scripts/")
        return

    for txt_file in txt_files:
        print(f"\nProcessing file: {txt_file.name}")
        with txt_file.open("r", encoding="utf-8") as f:
            full_text = f.read().strip()

        if not full_text:
            print(f"Skipping empty file: {txt_file.name}")
            continue

        # Split text into <= 5000 byte chunks
        chunks = split_text_into_chunks(full_text, max_bytes=5000)
        total_chunks = len(chunks)
        print(f"  File split into {total_chunks} chunk(s)")

        for i, chunk in enumerate(chunks, start=1):
            audio_file = OUTPUT_DIR / f"{txt_file.stem}_part{i}.mp3"
            try:
                synthesize_text_chunk(chunk, audio_file, i, total_chunks)
            except Exception as e:
                print(f"Error processing {txt_file.name} part {i}: {e}")

        print(f"Finished processing {txt_file.name} ({total_chunks} chunk(s) generated)")

    print("\nAll files processed successfully.")


# ============================================================
# SECTION: titles
# ============================================================

openai.api_key = os.getenv("OPENAI_API_KEY")

INPUT_DIR_TITLES = Path("output_scripts")
OUTPUT_FILE_TITLES = Path("chapter_titles.txt")

def get_first_paragraph(text):
    paragraphs = [p.strip() for p in text.split("\n\n") if p.strip()]
    return paragraphs[0] if paragraphs else text.strip()

def generate_chapter_title(first_paragraph):
    prompt = (
        "You are a helpful assistant. Based on the following text, "
        "provide a concise, human-friendly chapter title (5 words max):\n\n"
        f"{first_paragraph}"
    )
    response = openai_client.chat.completions.create(
        model="gpt-5.2",
        messages=[{"role": "user", "content": prompt}],
        temperature=0.5,
        max_completion_tokens=20
    )
    return response.choices[0].message.content.strip()


def run_titles():
    chapter_titles = {}
    txt_files = sorted(INPUT_DIR_TITLES.glob("*.txt"))
    if not txt_files:
        print("No text files found in output_scripts/")
        return

    for i, txt_file in enumerate(txt_files, start=1):
        with txt_file.open("r", encoding="utf-8") as f:
            text = f.read()

        first_paragraph = get_first_paragraph(text)
        title = generate_chapter_title(first_paragraph)

        if title:
            chapter_titles[txt_file.stem] = title
            print(f"{txt_file.stem}: {title}")
        else:
            chapter_titles[txt_file.stem] = f"Chapter {i}"
            print(f"{txt_file.stem}: Chapter {i} (default)")

    # Save titles to file
    with OUTPUT_FILE_TITLES.open("w", encoding="utf-8") as f:
        for key, title in chapter_titles.items():
            f.write(f"{key}: {title}\n")

    print(f"\nGenerated chapter titles saved to {OUTPUT_FILE_TITLES}")





# ============================================================
# SECTION: merge into a .m4b
# ============================================================

MERGED_DIR = Path("merged_audio")
AUDIO_INPUT = Path("audio_output")
CHAPTER_TITLES_FILE = Path("chapter_titles.txt")
METADATA_FILE = Path("chapters.ffmeta")
OUTPUT_FILE_M4B = Path("Course_Audiobook.m4b")

# -----------------------------
# Helper: duration in ms
# -----------------------------
def duration_ms(path: Path) -> int:
    r = subprocess.run(
        ["ffprobe", "-v", "error",
         "-show_entries", "format=duration",
         "-of", "default=nokey=1:noprint_wrappers=1",
         str(path)],
        capture_output=True,
        text=True,
        check=True
    )
    return int(float(r.stdout.strip()) * 1000)


def run_merge():
    MERGED_DIR.mkdir(exist_ok=True)

    # -----------------------------
    # Load chapter titles
    # -----------------------------
    chapter_titles = {}
    with CHAPTER_TITLES_FILE.open("r", encoding="utf-8") as f:
       for line in f:
           if ":" in line:
               key, title = line.strip().split(":", 1)
               chapter_titles[key.strip()] = title.strip()

    # -----------------------------
    # Group chunked MP3s
    # -----------------------------
    chunks = {}
    for mp3 in AUDIO_INPUT.glob("*.mp3"):
        match = re.match(r"(.*)_part(\d+)\.mp3", mp3.name)
        if match:
            chapter, part = match.groups()
            chunks.setdefault(chapter, []).append((int(part), mp3))

    if not chunks:
       raise RuntimeError("No chunked MP3 files found in audio_output/")

    # -----------------------------
    # Merge chunks with FFmpeg
    # -----------------------------
    merged_files = []

    for chapter, parts in sorted(chunks.items()):
        parts_sorted = sorted(parts)
        concat_file = MERGED_DIR / f"{chapter}_concat.txt"

        with concat_file.open("w", encoding="utf-8") as f:
            for _, mp3 in parts_sorted:
                f.write(f"file '{mp3.resolve()}'\n")

        output_mp3 = MERGED_DIR / f"{chapter}_full.mp3"

        subprocess.run([
            "ffmpeg",
            "-y",
            "-f", "concat",
            "-safe", "0",
            "-i", str(concat_file),
            "-c", "copy",
            str(output_mp3)
        ], check=True)

        merged_files.append(output_mp3)
        print(f"Merged chapter {chapter}")

    # -----------------------------
    # Build chapter metadata
    # -----------------------------
    lines = [";FFMETADATA1"]
    start = 0

    for mp3 in merged_files:
        chapter_id = mp3.stem.split("_")[0]
        title = chapter_titles.get(chapter_id, f"Chapter {chapter_id}")
        length = duration_ms(mp3)

        lines.extend([
            "[CHAPTER]",
            "TIMEBASE=1/1000",
            f"START={start}",
            f"END={start + length}",
            f"title={title}"
        ])

        start += length

    METADATA_FILE.write_text("\n".join(lines), encoding="utf-8")

    # -----------------------------
    # Final concat to M4B
    # -----------------------------
    concat_all = Path("concat_all.txt")
    with concat_all.open("w", encoding="utf-8") as f:
        for mp3 in merged_files:
            f.write(f"file '{mp3.resolve()}'\n")

    subprocess.run([
        "ffmpeg",
        "-f", "concat",
        "-safe", "0",
        "-i", str(concat_all),
        "-i", str(METADATA_FILE),
        "-map_metadata", "1",
        "-c:a", "aac",
        "-b:a", "128k",
        "-vn",
        str(OUTPUT_FILE_M4B)
    ], check=True)

    print(f"\nAudiobook created: {OUTPUT_FILE_M4B}")







# ============================================================
# SECTION: UNIFIED CLI + PIPELINE ORCHESTRATION
# ============================================================

def main():
    parser = argparse.ArgumentParser(
        description="Merged audio pipeline (zero-regression)"
    )

    parser.add_argument("--skip-chunk", action="store_true")
    parser.add_argument("--skip-generate", action="store_true")
    parser.add_argument("--skip-tts", action="store_true")
    parser.add_argument("--skip-titles", action="store_true")
    parser.add_argument("--skip-merge", action="store_true")

    args, unknown = parser.parse_known_args()

    if not args.skip_chunk:
        run_chunk(unknown)

    if not args.skip_generate:
        run_generate()

    if not args.skip_tts:
        run_tts()

    if not args.skip_titles:
        run_titles()

    if not args.skip_merge:
        run_merge()

if __name__ == "__main__":
    main()