#!/usr/bin/env python3
"""
Local curriculum outline generator.

Generates lesson outline JSON using a local llama.cpp model (Qwen3-Coder-30B).
Includes automatic validation + retry loop to handle schema errors.

Usage:
    # Generate a single week of History (5 lessons)
    python local_generate.py --subject History --week 1

    # Generate specific lessons
    python local_generate.py --subject ELA --weeks 22,23,24

    # Dry run — show prompt, don't call the model
    python local_generate.py --subject History --week 1 --dry-run
"""

import argparse
import json
import os
import sys
import time
import urllib.request
import urllib.error
from pathlib import Path
from typing import List, Optional

# Schema validation
sys.path.insert(0, str(Path(__file__).parent))
from curriculum_schema import validate_outline

# Model endpoint — Qwen3-Coder-30B on port 8082 (faster for structured output)
MODEL_URL = os.environ.get("LLAMA_URL", "http://localhost:8082/v1/chat/completions")
MODEL_NAME = os.environ.get("LLAMA_MODEL", "Qwen3-Coder-30B-A3B-Instruct-Q4_K_M.gguf")

# Retry config
MAX_RETRIES = 3
RETRY_DELAY = 1  # seconds between retries


def load_course_plan(subject: str) -> str:
    """Load the course plan markdown for context."""
    plan_path = Path(f"2nd_Grade/{subject}/32_week_course_plan.md")
    if plan_path.exists():
        return plan_path.read_text()
    return ""


def build_prompt(subject: str, week: int, day: str, topic: str, course_plan: str) -> str:
    """Build the system prompt for generating a SINGLE lesson."""
    
    # Subject-specific config
    config = {
        "History": {
            "days": ["Monday", "Tuesday", "Wednesday", "Thursday", "Friday"],
            "pages_per_lesson": 3,
            "font": "Georgia",
            "palette": "warm amber (#ff9800, #e65100, #fff3e0)",
        },
        "ELA": {
            "days": ["Monday", "Tuesday", "Wednesday", "Thursday"],
            "pages_per_lesson": 3,
            "font": "Comic Neue",
            "palette": "teal and coral (#00897b, #ff7043, #e0f2f1)",
        },
        "Science": {
            "days": ["Monday", "Tuesday", "Wednesday", "Thursday"],
            "pages_per_lesson": 3,
            "font": "Georgia",
            "palette": "green (#4caf50, #2e7d32, #e8f5e9)",
        },
    }

    cfg = config.get(subject, config["History"])
    
    # Extract week-specific context from course plan
    week_context = ""
    for line in course_plan.split("\n"):
        if f"Week {week}" in line or f"W{week}" in line:
            week_context += line + "\n"
        elif week_context and line.strip():
            week_context += line + "\n"
    week_context = week_context[:1000] if week_context else course_plan[:2000]
    
    prompt = f"""You are a curriculum designer generating a SINGLE lesson outline.

## Task

Generate ONE lesson in JSON format for:
- Subject: {subject}
- Week: {week}
- Day: {day}
- Topic: {topic}

## JSON Schema (output this ONE object):

{{
  "week": {week},
  "day": "{day}",
  "topic": "{topic}",
  "subtitle": "<friendly subtitle>",
  "icon": "<emoji>",
  "pages": [
    {{
      "title": "<page 1 title>",
      "sections": [
        {{
          "type": "reading",
          "title": "<section title>",
          "icon": "<emoji>",
          "paragraphs": [
            "A paragraph of text at 80-120 words. Write clearly for a 2nd grader.",
            "Another paragraph..."
          ]
        }},
        {{
          "type": "questions",
          "title": "<section title>",
          "questions": [
            {{"text": "Question text?"}},
            {{"text": "Another question?"}}
          ]
        }}
      ]
    }},
    {{
      "sections": [
        {{
          "type": "vocab_grid",
          "title": "Vocabulary",
          "words": [
            {{"term": "Word", "definition": "Simple definition"}}
          ]
        }},
        {{
          "type": "activity",
          "title": "Activity",
          "content": "Instructions for a hands-on activity."
        }}
      ]
    }}
  ]
}}

## Section types (pick a mix):
- **reading** — `paragraphs[]` (80-120 words each, 2nd grade level)
- **questions** — `questions[]` with `text`
- **vocab_grid** — `words[]` with `term` and `definition`
- **true_false** — `items[]` with `statement` and `answer`
- **fill_blank** — `items[]` with `sentence` and `answer`
- **sort_table** — `header` (2-3 cols) + `rows[]`
- **review_table** — `header` + `rows[]`
- **activity** — `content` (instructions)
- **think** — `questions[]` (discussion prompts)
- **timeline** — `entries[]` with `year` and `event`

## Rules:
1. Output ONLY valid JSON. No markdown, no explanation.
2. Generate {cfg['pages_per_lesson']} pages, each with 2-4 sections.
3. At least one reading section per lesson.
4. Include at least one active exercise (questions, activity, think, etc.).
5. Reading paragraphs: 80-120 words, 2nd grade reading level.
6. Factually accurate, engaging, age-appropriate.

## Week {week} context:
{week_context}

Generate the lesson JSON now:
"""
    return prompt


# Week topics — derived from course plan for each subject
WEEK_TOPICS = {
    "History": {
        1: ["What Is History", "Timelines", "Clues From The Past", "My Family History", "History Review"],
        2: ["Gift Of The Nile", "The Great Flood", "Boats And Tools", "Rivers Then & Now", "Nile Review"],
        3: ["Meet The Pharaohs", "Building Pyramids", "King Tut Treasures", "Leaders Then & Now", "Pharaohs Review"],
        4: ["Secret Picture Writing", "Rosetta Stone", "Mummies", "Writing Then & Now", "Hieroglyphics Review"],
        5: ["Life By The Nile", "Egyptian Jobs", "Games And Fun", "Kids Day Then & Now", "Egypt Grand Review"],
        6: ["Welcome To Greece", "City States", "Life By The Sea", "Geography And You", "Greece Review"],
        7: ["Gods Of Olympus", "Heroes And Monsters", "Trojan Horse", "Myths Today", "Myths Review"],
        8: ["First Olympics", "Olympic Events", "Day At Olympia", "Olympics Then & Now", "Olympics Review"],
        9: ["Gifts From Greece", "Democracy", "Greek Theater", "Great Thinkers", "Greek Ideas Review"],
        10: ["Welcome To Rome", "Emperors And Army", "Roman City Life", "Rome In Our World", "Roman Empire Review"],
        11: ["Master Builders", "Roman Roads", "Amazing Aqueducts", "The Colosseum", "Engineering Review"],
        12: ["Castles And Knights", "Inside A Castle", "Coat Of Arms", "Knights Code Today", "Castles Review"],
        13: ["Who Were Vikings", "Viking Longships", "Runes And Crafts", "Vikings Today", "Vikings Review"],
        14: ["Feudal System", "Village Life", "Market Day", "Then And Now", "Medieval Review"],
        15: ["Medieval Inventions", "Printing Press", "Be An Inventor", "Inventions Around You", "Inventions Review"],
        16: ["Age Of Exploration", "Ships And Tools", "Be A Navigator", "You Are An Explorer", "Exploration Review"],
        17: ["New Worlds Met", "Great Exchange", "Trading Without Words", "Two Sides Of A Story", "New Worlds Review"],
        18: ["Many Nations", "Journey Through Regions", "Culture Map", "Whose Land", "Diversity Review"],
        19: ["Native Homes", "Three Sisters", "Tools And Inventions", "Foods On Your Table", "Homes Food Tools Review"],
        20: ["Power Of Stories", "Legends And Tricksters", "Be A Storyteller", "Your Family Traditions", "Stories Review"],
        21: ["What Makes A Leader", "Sacagawea And Sequoyah", "Invent Your Symbols", "Standing Up", "Leaders Review"],
        22: ["Nations Today", "Code Talkers", "Scavenger Hunt", "Honoring And Celebrating", "Living History Review"],
        23: ["First Colonies", "Jamestown", "Mayflower", "Why People Move", "Colonies Review"],
        24: ["Colonial Life", "Colonial Towns", "Colonial School", "Colonial Kids", "Colonial Review"],
        25: ["Road To Revolution", "Boston Tea Party", "War For Independence", "Freedom Today", "Revolution Review"],
        26: ["Founding Fathers", "The Declaration", "Ben Franklin", "Rules And Rights", "Founders Review"],
        27: ["Early Presidents", "George Washington", "Thomas Jefferson", "Voting And Me", "Presidents Review"],
        28: ["Growing Nation", "Oregon Trail", "Pack Your Wagon", "Journeys Then Now", "Westward Review"],
        29: ["Civil War Begins", "Life During War", "War Timeline", "Solving Disagreements", "Civil War Review"],
        30: ["Young Lincoln", "Ending Slavery", "Lincolns Words", "Honesty Like Abe", "Lincoln Review"],
        31: ["Amazing 1900s", "Planes And Cars", "Invention Lab", "Inventions At Home", "1900s Review"],
        32: ["America Today", "Freedoms And Symbols", "My History Museum", "Part Of History", "Year In Review"],
    },
}

DAYS = ["Monday", "Tuesday", "Wednesday", "Thursday", "Friday"]


def call_model(prompt: str, max_tokens: int = 8000, temperature: float = 0.7) -> str:
    """Call the local model via OpenAI-compatible API."""
    messages = [
        {"role": "system", "content": "You are a curriculum designer. Output ONLY valid JSON. No markdown, no explanation."},
        {"role": "user", "content": prompt}
    ]
    
    data = json.dumps({
        "model": MODEL_NAME,
        "messages": messages,
        "max_tokens": max_tokens,
        "temperature": temperature,
    }).encode()

    req = urllib.request.Request(
        MODEL_URL,
        data=data,
        headers={"Content-Type": "application/json"},
    )
    
    try:
        with urllib.request.urlopen(req, timeout=300) as resp:
            result = json.loads(resp.read().decode())
            return result["choices"][0]["message"]["content"]
    except urllib.error.URLError as e:
        raise ConnectionError(f"Model API error: {e}")
    except Exception as e:
        raise RuntimeError(f"Unexpected error calling model: {e}")


def extract_json(text: str) -> list:
    """Extract JSON array from model output, handling markdown fences."""
    # Strip markdown code fences
    text = text.strip()
    if text.startswith("```"):
        # Remove ```json ... ``` fences
        lines = text.split("\n")
        if lines[0].startswith("```"):
            lines = lines[1:]
        if lines and lines[-1].strip() == "```":
            lines = lines[:-1]
        text = "\n".join(lines)
    
    # Try to find JSON array
    start = text.find("[")
    end = text.rfind("]")
    if start == -1 or end == -1:
        raise ValueError("No JSON array found in output")
    
    return json.loads(text[start:end + 1])


def extract_json_single(text: str):
    """Extract a single JSON object from model output."""
    text = text.strip()
    if text.startswith("```"):
        lines = text.split("\n")
        if lines[0].startswith("```"):
            lines = lines[1:]
        if lines and lines[-1].strip() == "```":
            lines = lines[:-1]
        text = "\n".join(lines)
    
    start = text.find("{")
    end = text.rfind("}")
    if start == -1 or end == -1:
        raise ValueError("No JSON object found in output")
    
    return json.loads(text[start:end + 1])


def validate_and_fix(lessons: list, subject: str) -> List[tuple]:
    """Validate each lesson, return list of (valid, lesson, errors)."""
    results = []
    for i, lesson in enumerate(lessons):
        valid, errors = validate_outline(lesson)
        results.append((valid, lesson, errors))
    return results


def save_lesson(lesson: dict, output_dir: str, subject: str) -> str:
    """Save a lesson outline to JSON. Returns filepath."""
    week = lesson["week"]
    day = lesson["day"]
    topic = lesson["topic"].replace(" ", "_")
    
    filename = f"W{week:02d}_{day}_{topic}.json"
    path = os.path.join(output_dir, filename)
    
    with open(path, "w") as f:
        json.dump(lesson, f, indent=2, ensure_ascii=False)
    
    return path


def generate_week(
    subject: str,
    week: int,
    output_dir: str,
    course_plan: str,
    dry_run: bool = False,
) -> List[str]:
    """Generate a full week of lessons, one at a time."""
    topics = WEEK_TOPICS.get(subject, {}).get(week)
    if not topics:
        print(f"No topics defined for {subject} Week {week}")
        return []
    
    days = DAYS[:len(topics)]
    
    if dry_run:
        print("=== DRY RUN ===")
        prompt = build_prompt(subject, week, days[0], topics[0], course_plan)
        print(prompt[:2000])
        print("...")
        return []
    
    print(f"Generating Week {week} ({subject}) — {len(days)} lessons")
    
    saved = []
    for i, (day, topic) in enumerate(zip(days, topics)):
        print(f"\n  [{i+1}/{len(days)}] {day}: {topic}")
        
        # Check if already exists
        existing = [f for f in os.listdir(output_dir)
                    if f.startswith(f"W{week:02d}_{day}_") and f.endswith(".json")]
        if existing:
            print(f"    Skip — {existing[0]} exists")
            saved.append(os.path.join(output_dir, existing[0]))
            continue
        
        prompt = build_prompt(subject, week, day, topic, course_plan)
        
        lesson = None
        for attempt in range(MAX_RETRIES):
            if attempt > 0:
                print(f"    Retry {attempt}/{MAX_RETRIES}...")
                time.sleep(RETRY_DELAY)
            
            try:
                raw = call_model(prompt, temperature=0.7)
                lesson = extract_json_single(raw)
            except (json.JSONDecodeError, ValueError, KeyError, IndexError) as e:
                print(f"    Parse error (attempt {attempt+1}): {e}")
                if attempt == MAX_RETRIES - 1:
                    debug_path = os.path.join(output_dir, f"W{week:02d}_{day}_debug.txt")
                    with open(debug_path, "w") as f:
                        f.write(raw)
                    print(f"    Debug saved: {debug_path}")
                continue
            
            # Validate
            valid, errors = validate_outline(lesson)
            if valid:
                break
            else:
                print(f"    Schema error: {', '.join(errors[:3])}")
        
        if lesson and validate_outline(lesson)[0]:
            path = save_lesson(lesson, output_dir, subject)
            print(f"    ✓ Saved: {os.path.basename(path)}")
            saved.append(path)
        else:
            print(f"    ✗ Failed after {MAX_RETRIES} attempts")
    
    return saved


def main():
    parser = argparse.ArgumentParser(description="Generate curriculum outlines locally")
    parser.add_argument("--subject", choices=["History", "Science", "ELA"], default="History")
    parser.add_argument("--week", type=int, help="Generate a single week")
    parser.add_argument("--weeks", type=str, help="Generate multiple weeks (comma-separated)")
    parser.add_argument("--output-dir", type=str, help="Output directory for JSON files")
    parser.add_argument("--dry-run", action="store_true", help="Show prompt without calling model")
    parser.add_argument("--max-tokens", type=int, default=8000, help="Max tokens per request")
    
    args = parser.parse_args()
    
    subject = args.subject
    
    # Set output dir
    if args.output_dir:
        output_dir = args.output_dir
    else:
        output_dir = f"2nd_Grade/{subject}/outlines"
    
    os.makedirs(output_dir, exist_ok=True)
    
    # Load course plan
    course_plan = load_course_plan(subject)
    
    # Determine weeks
    if args.week:
        weeks = [args.week]
    elif args.weeks:
        weeks = [int(w.strip()) for w in args.weeks.split(",")]
    else:
        print("Error: specify --week or --weeks")
        sys.exit(1)
    
    # Generate
    total_saved = 0
    for w in weeks:
        print(f"\n{'='*60}")
        print(f"Subject: {subject}, Week: {w}")
        print(f"{'='*60}")
        
        # Check existing files
        existing = [f for f in os.listdir(output_dir) if f.startswith(f"W{w:02d}_")]
        if existing:
            print(f"Found {len(existing)} existing files for week {w}")
        
        paths = generate_week(
            subject=subject,
            week=w,
            output_dir=output_dir,
            course_plan=course_plan,
            dry_run=args.dry_run,
        )
        total_saved += len(paths)
    
    print(f"\n{'='*60}")
    print(f"Done! Saved {total_saved} lesson(s)")
    
    # Validate all output
    all_files = sorted(Path(output_dir).glob(f"W*_*.json"))
    print(f"\nValidating {len(all_files)} files...")
    passed = 0
    for f in all_files:
        with open(f) as fh:
            lesson = json.load(fh)
        valid, _ = validate_outline(lesson)
        if valid:
            passed += 1
    print(f"  {passed}/{len(all_files)} valid ({100*passed//len(all_files)}%)")


if __name__ == "__main__":
    main()