run_teachcraft.py

"""
TeachCraft Runner - Generate lesson from source documents + requirement

Pipeline:
1. Load source documents
2. Ingest into Knowledge Pipeline
3. Agentic Exploration (Intent Parsing + Document Exploration)
4. Generate Lesson (Plan → Outline → Pages)
5. Output in LessonBench format (slides + scripts + quiz)
"""

import asyncio
import json
import logging
from pathlib import Path
from typing import Dict, List, Optional

from data.loader import load_case, get_source_docs_dir
from generator.pipeline import CourseGenerationPipeline
from generator.script import generate_scripts_batch
from knowledge.pipeline import KnowledgePipeline
from knowledge.models import DocumentImage
from exploration import explore_documents, ExplorationResult
from models.pipeline_schema import PageType
from utils.debug import set_debug_mode

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)


async def render_canvas_to_image(
    canvas_json: dict,
    output_path: Path,
    image_id_map: Optional[Dict[str, str]] = None
) -> bool:
    """Render canvas JSON to PNG image

    Args:
        canvas_json: Canvas JSON data
        output_path: Output path for the PNG image
        image_id_map: Mapping from image IDs to base64 data URLs
    """
    try:
        from renderer.playwright_renderer import PPTistPlaywrightRenderer

        renderer = PPTistPlaywrightRenderer()
        await renderer.initialize()

        result = await renderer.render_canvas_component(canvas_json, image_id_map=image_id_map)

        with open(output_path, 'wb') as f:
            f.write(result['image_data'])

        # Cleanup browser resources
        if hasattr(renderer, 'close'):
            await renderer.close()

        return True
    except Exception as e:
        logger.error(f"Failed to render canvas: {e}")
        return False


def format_exploration_result(result: ExplorationResult) -> tuple[str, List[DocumentImage]]:
    """Format exploration result as context for lesson generation.

    Returns:
        Tuple of (text_context, collected_images)
    """
    chunks_by_intent = result.chunk_pool.chunks_by_intent

    sections = []
    all_images: List[DocumentImage] = []
    seen_uris = set()  # Deduplicate images

    for intent, chunk_ids in chunks_by_intent.items():
        if not chunk_ids:
            continue
        section_text = f"## {intent} Content\n\n"
        # Get actual chunk objects from IDs
        chunks = result.chunk_pool.get_chunks_by_ids(chunk_ids[:5])  # Limit chunks per intent
        for chunk in chunks:
            section_text += f"### From: {chunk.provenance.source_document_id}\n"
            section_text += f"{chunk.text}\n\n"

            # Collect images from chunk
            if chunk.images:
                for img_data in chunk.images:
                    uri = img_data.get("uri", "")
                    if uri and uri not in seen_uris:
                        seen_uris.add(uri)
                        all_images.append(DocumentImage(
                            uri=uri,
                            caption=img_data.get("caption"),
                            page_number=img_data.get("page_number"),
                            width=img_data.get("width"),
                            height=img_data.get("height"),
                            format=img_data.get("format"),
                        ))

        sections.append(section_text)

    text_context = "\n".join(sections)
    logger.info(f"Collected {len(all_images)} unique images from exploration")
    return text_context, all_images


async def run_teachcraft(
    case_id: str,
    model: str = "gpt-4o",
    exploration_model: str = "claude-haiku-4-5",
    debug: bool = True,
    skip_exploration: bool = False,
    skip_planning: bool = False,
    latex_output: bool = False,
    output_dir: Optional[Path] = None,
):
    """
    Run TeachCraft pipeline on a LessonBench case

    Args:
        case_id: LessonBench case ID
        model: Model for generation
        exploration_model: Model for exploration agents
        debug: Enable debug mode
        skip_exploration: Skip agentic exploration (use raw docs) - for teachcraft_no_exp ablation
        skip_planning: Skip hierarchical planning (direct generation) - for teachcraft_no_plan ablation
        latex_output: Generate LaTeX Beamer instead of Canvas JSON - for teachcraft_no_schema ablation
        output_dir: Override output directory (default: outputs/{case_id}/{model})
    """
    # Enable debug mode
    if debug:
        set_debug_mode(True)

    # Load case
    case = load_case(case_id)
    logger.info(f"Loaded case: {case.case_id}")
    logger.info(f"Requirement: {case.requirement}")

    # Get source document paths
    source_dir = get_source_docs_dir(case_id)
    doc_paths = [str(source_dir / doc) for doc in case.documents]
    doc_ids = case.documents

    logger.info(f"Source documents: {doc_ids}")

    # Output directory - use provided or default to outputs/{case_id}/{model}
    if output_dir is None:
        model_short_name = model.split("/")[-1] if "/" in model else model  # Handle provider/model format
        output_dir = Path(f"outputs/{case_id}/{model_short_name}")
    else:
        output_dir = Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    logger.info(f"Output directory: {output_dir}")

    # Step 1: Agentic Exploration (or skip)
    exploration_images: List[DocumentImage] = []  # Initialize empty

    # Helper function to load standalone image files
    def load_standalone_images(doc_paths: List[str], doc_ids: List[str]) -> List[DocumentImage]:
        """Load standalone image files (png, jpg, etc.) as DocumentImage objects"""
        from PIL import Image as PILImage
        images = []
        for doc_path, doc_id in zip(doc_paths, doc_ids):
            suffix = Path(doc_path).suffix.lower()
            if suffix in ['.png', '.jpg', '.jpeg', '.gif', '.webp']:
                try:
                    with PILImage.open(doc_path) as pil_img:
                        width, height = pil_img.size
                        img_format = pil_img.format.lower() if pil_img.format else suffix[1:]

                    images.append(DocumentImage(
                        uri=str(Path(doc_path).absolute()),
                        caption=doc_id,
                        page_number=None,
                        width=width,
                        height=height,
                        format=img_format,
                    ))
                    logger.info(f"Loaded standalone image: {doc_id} ({width}x{height})")
                except Exception as e:
                    logger.warning(f"Failed to load image {doc_id}: {e}")
        return images

    # Load standalone image files from source docs
    exploration_images = load_standalone_images(doc_paths, doc_ids)

    if skip_exploration:
        logger.info("Skipping agentic exploration, using raw document content...")
        # Fallback: load documents directly (with cache support)
        from knowledge.processor import DoclingProcessor
        from knowledge.video_processor import VideoProcessor, VIDEO_EXTENSIONS
        from knowledge.cache import load_cached_document

        processor = None  # Lazy initialization
        video_processor = None

        doc_contents = []
        for doc_path, doc_id in zip(doc_paths, doc_ids):
            suffix = Path(doc_path).suffix.lower()

            if suffix in ['.png', '.jpg', '.jpeg', '.gif', '.webp']:
                doc_contents.append(f"## {doc_id}\n[Image document - see attached]\n")

            elif suffix in VIDEO_EXTENSIONS:
                # Videos: check cache first
                cached_doc = load_cached_document(case_id, doc_id)
                if cached_doc:
                    logger.info(f"Loaded video from cache: {doc_id}")
                    doc_contents.append(f"## {doc_id}\n{cached_doc.content}\n")
                    exploration_images.extend(cached_doc.images)
                else:
                    # 使用 VideoProcessor 处理视频（提取关键帧+转录）
                    try:
                        if processor is None:
                            processor = DoclingProcessor(enable_asr=True)
                        if video_processor is None:
                            video_processor = VideoProcessor(
                                processor=processor,
                                frame_interval=15,
                                max_frames=20
                            )
                        video_doc = video_processor.process(doc_path, doc_id=doc_id)
                        doc_contents.append(f"## {doc_id}\n{video_doc.content}\n")
                        exploration_images.extend(video_doc.images)
                        logger.info(f"Processed video: {doc_id}, extracted {len(video_doc.images)} keyframes")
                    except Exception as e:
                        logger.warning(f"Failed to process video {doc_id}: {e}")

            elif suffix == '.url':
                # .url files: check cache first
                cached_doc = load_cached_document(case_id, doc_id)
                if cached_doc:
                    logger.info(f"Loaded URL doc from cache: {doc_id}")
                    doc_contents.append(f"## {doc_id}\n{cached_doc.content}\n")
                    exploration_images.extend(cached_doc.images)
                else:
                    try:
                        if processor is None:
                            processor = DoclingProcessor(enable_asr=True)
                        with open(doc_path, 'r') as f:
                            url = f.read().strip()
                        logger.info(f"Processing URL from {doc_id}: {url}")
                        doc = processor.process(url, doc_id=doc_id)
                        doc_contents.append(f"## {doc_id}\nSource: {url}\n\n{doc.content}\n")
                        exploration_images.extend(doc.images)
                        logger.info(f"Processed URL: {doc_id}, {len(doc.content)} chars, {len(doc.images)} images")
                    except Exception as e:
                        logger.warning(f"Failed to process URL {doc_id}: {e}")

            elif Path(doc_path).exists():
                # Regular documents: check cache first
                cached_doc = load_cached_document(case_id, doc_id)
                if cached_doc:
                    logger.info(f"Loaded from cache: {doc_id}")
                    doc_contents.append(f"## {doc_id}\n{cached_doc.content}\n")
                    exploration_images.extend(cached_doc.images)
                else:
                    try:
                        if processor is None:
                            processor = DoclingProcessor(enable_asr=True)
                        doc = processor.process(doc_path, doc_id=doc_id)
                        doc_contents.append(f"## {doc_id}\n{doc.content}\n")
                        exploration_images.extend(doc.images)
                    except Exception as e:
                        logger.warning(f"Failed to process {doc_id}: {e}")
            else:
                logger.warning(f"Document not found: {doc_path}")

        exploration_context = "\n---\n".join(doc_contents)
        logger.info(f"Collected {len(exploration_images)} total images")
    else:
        logger.info("Starting agentic exploration...")

        # Initialize Knowledge Pipeline
        knowledge_pipeline = KnowledgePipeline()

        # Ingest documents (with cache support)
        for doc_path, doc_id in zip(doc_paths, doc_ids):
            suffix = Path(doc_path).suffix.lower()
            if suffix in ['.png', '.jpg', '.jpeg']:
                logger.info(f"Skipping image: {doc_id}")
                continue
            logger.info(f"Ingesting: {doc_id}")
            try:
                if suffix == '.url':
                    # .url 文件：读取 URL 并 ingest
                    with open(doc_path, 'r') as f:
                        url = f.read().strip()
                    logger.info(f"Ingesting URL from {doc_id}: {url}")
                    knowledge_pipeline.ingest(url, doc_id=doc_id, case_id=case_id)
                else:
                    knowledge_pipeline.ingest(doc_path, doc_id=doc_id, case_id=case_id)
            except Exception as e:
                logger.warning(f"Failed to ingest {doc_id}: {e}")

        # Run exploration
        exploration_result = await explore_documents(
            requirements=case.requirement,
            document_paths=[p for p, d in zip(doc_paths, doc_ids)
                          if not Path(p).suffix.lower() in ['.png', '.jpg', '.jpeg']],
            knowledge_pipeline=knowledge_pipeline,
            document_ids=[d for d, p in zip(doc_ids, doc_paths)
                         if not Path(p).suffix.lower() in ['.png', '.jpg', '.jpeg']],
            agent_model=exploration_model,
        )

        logger.info(f"Exploration complete: {exploration_result.chunk_pool.total_chunks} chunks collected")

        # Save exploration result
        exploration_dump_path = output_dir / "exploration_result.json"
        with open(exploration_dump_path, 'w') as f:
            json.dump({
                "exploration_id": exploration_result.exploration_id,
                "requirements": exploration_result.requirements,
                "total_chunks": exploration_result.chunk_pool.total_chunks,
                "chunks_by_intent": {
                    k: len(v) for k, v in exploration_result.chunk_pool.chunks_by_intent.items()
                },
                "exploration_log": exploration_result.exploration_log,
            }, f, ensure_ascii=False, indent=2)
        logger.info(f"Exploration result saved to: {exploration_dump_path}")

        exploration_context, collected_images = format_exploration_result(exploration_result)
        exploration_images.extend(collected_images)  # Extend, don't overwrite standalone images

        # Also collect images from ALL document chunks (not just selected ones)
        # This ensures we don't miss images from documents whose chunks weren't selected
        import json as json_module
        seen_uris = {img.uri for img in exploration_images}  # Track already added images
        for doc_id in doc_ids:
            if Path(doc_paths[doc_ids.index(doc_id)]).suffix.lower() in ['.png', '.jpg', '.jpeg']:
                continue
            try:
                all_chunks = knowledge_pipeline.store.get_document_chunks(doc_id)
                for chunk in all_chunks:
                    if "images" in chunk.metadata:
                        images_data = chunk.metadata["images"]
                        if isinstance(images_data, str):
                            images_data = json_module.loads(images_data)
                        for img_data in images_data:
                            uri = img_data.get("uri", "")
                            if uri and uri not in seen_uris:
                                seen_uris.add(uri)
                                exploration_images.append(DocumentImage(
                                    uri=uri,
                                    caption=img_data.get("caption"),
                                    page_number=img_data.get("page_number"),
                                    width=img_data.get("width"),
                                    height=img_data.get("height"),
                                    format=img_data.get("format"),
                                ))
                                logger.info(f"Added image from {doc_id}: {Path(uri).name}")
            except Exception as e:
                logger.warning(f"Failed to collect images from {doc_id}: {e}")
        logger.info(f"Total images after collecting from all docs: {len(exploration_images)}")

    # Step 2: Build augmented requirement
    augmented_input = f"""## User Requirement
{case.requirement}

## Explored Content from Source Documents
{exploration_context}

Please use the information from these source documents to create the lesson content.
"""

    # Step 3: Run generation pipeline
    logger.info(f"Passing {len(exploration_images)} images to generation pipeline")

    # Handle ablation modes
    if latex_output:
        # TeachCraft-NoSchemaMemory ablation: Generate LaTeX Beamer instead of Canvas JSON
        logger.info("Running in LaTeX output mode (teachcraft_no_schema ablation)")
        from generator.latex_generator import generate_latex_lesson
        result = await generate_latex_lesson(
            requirements=augmented_input,
            model=model,
            output_dir=output_dir,
            available_images=exploration_images,
        )
    elif skip_planning:
        # TeachCraft-NoPlanning ablation: Skip Stage 1+2, direct generation
        logger.info("Running in direct generation mode (teachcraft_no_plan ablation)")
        pipeline = CourseGenerationPipeline(
            model=model,
            output_dir=output_dir,
            available_images=exploration_images,
        )
        result = await pipeline.generate_direct(
            requirements=augmented_input,
            save_intermediates=True,
        )
    else:
        # Normal TeachCraft flow
        pipeline = CourseGenerationPipeline(
            model=model,
            output_dir=output_dir,
            available_images=exploration_images,
        )
        result = await pipeline.generate(
            requirements=augmented_input,
            save_intermediates=True,
        )

    logger.info(f"Generation complete. Converting to LessonBench format...")

    # Step 4: Convert to LessonBench format
    lessonbench_dir = output_dir / "lessonbench"
    lessonbench_dir.mkdir(exist_ok=True)

    # For latex_output mode, generate_latex_lesson() already created PNG files and scripts
    # in the lessonbench directory, so skip the canvas rendering and script generation
    if latex_output:
        logger.info("LaTeX output mode: Skipping canvas rendering (already done by generate_latex_lesson)")
        # Count slides and save metadata
        num_slides = len(list(lessonbench_dir.glob("*_slide.png")))
        metadata = {
            "case_id": case_id,
            "model": model,
            "system": "teachcraft_no_schema",
            "num_slides": num_slides,
            "ablation": "latex_output",
        }
        metadata_path = output_dir / "metadata.json"
        with open(metadata_path, 'w') as f:
            json.dump(metadata, f, indent=2, ensure_ascii=False)
        logger.info(f"Saved metadata: {metadata_path}")
        logger.info(f"Generation complete: {num_slides} slides generated")
        return result

    # Build image_id_map for rendering
    import base64
    image_id_map: Dict[str, str] = {}
    for i, img in enumerate(exploration_images):
        try:
            img_path = img.uri
            if img_path.startswith("file://"):
                img_path = img_path[7:]
            with open(img_path, 'rb') as f:
                data = base64.b64encode(f.read()).decode()
            img_format = img.format or "jpeg"
            image_id_map[f"image_{i}"] = f"data:image/{img_format};base64,{data}"
            logger.debug(f"Built image_id_map entry: image_{i}")
        except Exception as e:
            logger.warning(f"Failed to load image for rendering: {img.uri}: {e}")

    if image_id_map:
        logger.info(f"Built image_id_map with {len(image_id_map)} images for rendering")

    # Step 4a: Render all canvas pages to images first
    canvas_pages = []  # List of (page, component_index, canvas_path, image_path)
    component_index = 0

    for page in result.pages:
        if page.page_type == PageType.canvas:
            # Save canvas JSON
            canvas_path = lessonbench_dir / f"{component_index:02d}_canvas.json"
            with open(canvas_path, 'w') as f:
                json.dump(page.content, f, ensure_ascii=False, indent=2)

            # Render to image
            image_path = lessonbench_dir / f"{component_index:02d}_slide.png"
            rendered = await render_canvas_to_image(page.content, image_path, image_id_map=image_id_map)
            if rendered:
                logger.info(f"Rendered slide: {image_path.name}")
            else:
                logger.warning(f"Failed to render slide {component_index}")

            canvas_pages.append((page, component_index, canvas_path, image_path if rendered else None))

        elif page.page_type == PageType.quiz:
            # Save quiz JSON
            quiz_path = lessonbench_dir / f"{component_index:02d}_quiz.json"
            with open(quiz_path, 'w') as f:
                json.dump(page.content, f, ensure_ascii=False, indent=2)
            logger.info(f"Saved quiz: {quiz_path.name}")

        component_index += 1

    # Step 4b: Batch generate scripts for all canvas pages
    if canvas_pages:
        logger.info(f"Batch generating scripts for {len(canvas_pages)} slides...")

        scripts = await generate_scripts_batch(
            slides=[p.content for p, _, _, _ in canvas_pages],
            lesson_plan=result.lesson_plan.content,
            slide_images=[img_path for _, _, _, img_path in canvas_pages],
            model=model,
        )

        # Save scripts
        for i, (page, comp_idx, canvas_path, image_path) in enumerate(canvas_pages):
            slide_num = i + 1  # Scripts are 1-indexed
            script_content = scripts.get(slide_num, f"[Script generation failed for slide {slide_num}]")

            script_path = lessonbench_dir / f"{comp_idx:02d}_script.md"
            script_path.write_text(script_content, encoding='utf-8')
            logger.info(f"Saved script: {script_path.name}")

    logger.info(f"LessonBench format output saved to: {lessonbench_dir}")
    return result


if __name__ == "__main__":
    import sys

    case_id = sys.argv[1] if len(sys.argv) > 1 else "algebra"
    model = sys.argv[2] if len(sys.argv) > 2 else "gpt-4o"

    # Parse flags
    debug = "--no-debug" not in sys.argv
    skip_exploration = "--skip-exploration" in sys.argv
    skip_planning = "--skip-planning" in sys.argv
    latex_output = "--latex-output" in sys.argv

    asyncio.run(run_teachcraft(
        case_id=case_id,
        model=model,
        debug=debug,
        skip_exploration=skip_exploration,
        skip_planning=skip_planning,
        latex_output=latex_output,
    ))