#!/usr/bin/env python3
"""
Treatment Reader Compiler
Compiles treatment.md into a clean, readable Fountain document for human review.

Usage:
    python compile_treatment.py <project_path>
    python compile_treatment.py <project_path> --output filename.fountain

Output: A clean Fountain file with just episode number, title, and prose paragraph.
Optimized for ~15 minute read-through. No emojis, no trackers, just the story.
"""

import re
import sys
from pathlib import Path
from datetime import date


def parse_treatment(filepath):
    """Parse treatment.md and extract episode data."""
    with open(filepath, 'r', encoding='utf-8') as f:
        content = f.read()

    # Extract series title
    title_match = re.search(r'^#\s*(.+?)\s*[-—]', content, re.MULTILINE)
    series_title = title_match.group(1).strip() if title_match else "UNTITLED"

    # Extract logline from the blockquote after title
    logline_match = re.search(r'^>\s*(.+?)$', content, re.MULTILINE)
    logline = logline_match.group(1).strip() if logline_match else ""

    episodes = []

    # Episode header pattern
    episode_pattern = re.compile(
        r'###\s*Episode\s*(\d+):\s*"([^"]+)"',
        re.IGNORECASE
    )

    # THE MOMENT pattern
    moment_pattern = re.compile(
        r'\*\*THE MOMENT:\*\*\s*(.+?)(?:\n|$)',
        re.IGNORECASE
    )

    # Cliffhanger image pattern
    cliffhanger_pattern = re.compile(
        r'\*\*\[CLIFFHANGER:\s*(.+?)\]\*\*',
        re.IGNORECASE
    )

    # Split into episode blocks
    episode_blocks = re.split(r'(?=###\s*Episode\s*\d+)', content)

    for block in episode_blocks:
        header_match = episode_pattern.search(block)
        if not header_match:
            continue

        ep_num = int(header_match.group(1))
        title = header_match.group(2)

        # Get THE MOMENT
        moment_match = moment_pattern.search(block)
        the_moment = moment_match.group(1).strip() if moment_match else ""

        # Get cliffhanger image
        cliff_match = cliffhanger_pattern.search(block)
        cliffhanger = cliff_match.group(1).strip() if cliff_match else ""

        # Extract prose (between THE MOMENT line and cliffhanger)
        prose_lines = []
        lines = block.split('\n')
        in_prose = False

        for line in lines:
            if '**THE MOMENT:**' in line:
                in_prose = True
                continue
            if '**[CLIFFHANGER:' in line:
                in_prose = False
                continue
            if in_prose and line.strip() and not line.startswith('**') and not line.startswith('---'):
                prose_lines.append(line.strip())

        prose = ' '.join(prose_lines)

        episodes.append({
            'episode': ep_num,
            'title': title,
            'the_moment': the_moment,
            'prose': prose,
            'cliffhanger': cliffhanger
        })

    return {
        'title': series_title,
        'logline': logline,
        'episodes': sorted(episodes, key=lambda x: x['episode'])
    }


def format_episode_fountain(ep):
    """Format a single episode for Fountain output."""
    lines = []

    # Episode header as section
    lines.append(f"[[EPISODE {ep['episode']}: {ep['title'].upper()}]]")
    lines.append("")

    # Prose paragraph as action block
    if ep['prose']:
        lines.append(ep['prose'])
        lines.append("")

    # Cliffhanger as the final beat
    if ep['cliffhanger']:
        lines.append(f"[{ep['cliffhanger']}]")
        lines.append("")

    return '\n'.join(lines)


def compile_treatment(project_path, output_name=None):
    """Compile treatment into clean Fountain document."""
    project = Path(project_path).resolve()
    treatment_path = project / 'treatment.md'

    if not treatment_path.exists():
        raise FileNotFoundError(f"treatment.md not found at {treatment_path}")

    data = parse_treatment(treatment_path)

    # Build Fountain output
    lines = []

    # Title page
    lines.append(f"Title: {data['title']} — Treatment")
    lines.append("Credit: A Microdrama Series")
    lines.append(f"Draft date: {date.today().isoformat()}")
    lines.append("")

    # Logline as synopsis
    if data['logline']:
        lines.append(f"Notes: {data['logline']}")
        lines.append("")

    # Page break after title
    lines.append("===")
    lines.append("")

    # Episodes
    for ep in data['episodes']:
        lines.append(format_episode_fountain(ep))

        # Page break between episodes
        lines.append("===")
        lines.append("")

    # Determine output path
    if output_name:
        if not output_name.endswith('.fountain'):
            output_name += '.fountain'
        output_path = project / output_name
    else:
        output_path = project / f"{data['title'].lower().replace(' ', '_')}_treatment.fountain"

    # Write file
    with open(output_path, 'w', encoding='utf-8') as f:
        f.write('\n'.join(lines))

    print(f"Treatment compiled to Fountain:")
    print(f"  {output_path}")
    print(f"  {len(data['episodes'])} episodes")

    # Word count estimate
    total_words = sum(len(ep['prose'].split()) for ep in data['episodes'])
    print(f"  ~{total_words:,} words (~{total_words // 200} min read)")

    return str(output_path)


def main():
    if len(sys.argv) < 2:
        print(__doc__)
        sys.exit(1)

    project_path = sys.argv[1]
    output_name = None

    # Parse --output flag
    if '--output' in sys.argv:
        idx = sys.argv.index('--output')
        if idx + 1 < len(sys.argv):
            output_name = sys.argv[idx + 1]

    try:
        compile_treatment(project_path, output_name)
    except Exception as e:
        print(f"Error: {e}")
        sys.exit(1)


if __name__ == '__main__':
    main()