"""
Generate markdown lessons based on extracted topics.
"""
import os
import json
from typing import Dict, Any, List, Optional

from slugify import slugify

from utils.io import save_output, save_markdown_lesson
from chains.base import build_chain, parse_output
from chains.lessons.extract_topics import load_topics

def parse_lesson(output: str) -> Dict[str, Any]:
    """
    Parse the generated lesson from the LLM output.

    Args:
        output: Raw output from the LLM

    Returns:
        Dict: Parsed lesson content
    """
    try:
        # Try to extract JSON
        import re
        json_blocks = re.findall(r'```(?:json)?\s*([\s\S]*?)```', output)

        if json_blocks:
            # Found JSON block
            try:
                return json.loads(json_blocks[0])
            except json.JSONDecodeError:
                pass

        # If no JSON blocks or parsing failed, try to find markdown content
        markdown_pattern = r'(?:```markdown\s*([\s\S]*?)```)|(?:#{1,6}.*[\s\S]*)'
        match = re.search(markdown_pattern, output)

        if match and match.group(1):
            # Found markdown in code block
            return {"content": match.group(1).strip()}
        elif "# " in output:
            # Found markdown without code block
            return {"content": output.strip()}

        # If all else fails, just return the raw output
        return {"content": output.strip()}
    except Exception as e:
        print(f"Error parsing lesson: {e}")
        # Return raw output as fallback
        return {"content": output.strip()}

def generate_lesson_for_topic(
    topic: Dict[str, str],
    run_id: str,
    index: int
) -> Dict[str, Any]:
    """
    Generate a markdown lesson for a single topic.

    Args:
        topic: Topic dictionary with 'topic', 'content', 'duration', 'notes'
        run_id: Run identifier
        index: Topic index for ordering

    Returns:
        Dict: Generated lesson with metadata
    """
    print(f"Generating lesson for topic: {topic['topic']}")

    # Prepare input variables for the chain
    input_variables = {
        "topic": topic["topic"],
        "content": topic["content"],
        "duration": topic["duration"],
        "notes": topic["notes"]
    }

    # Run the chain to generate the lesson
    result = build_chain(
        chain_name="generate_lessons",
        pipeline="lessons",
        run_id=run_id,
        input_variables=input_variables
    )

    # Parse the lesson content
    lesson = parse_lesson(result["output"])

    # Add metadata
    lesson["topic"] = topic["topic"]
    lesson["duration"] = topic["duration"]
    lesson["index"] = index

    # Save the lesson to a JSON file
    save_output(
        data=lesson,
        pipeline="lessons",
        step=f"lesson_{index:02d}",
        run_id=run_id,
        subfolder="lessons"
    )

    # Save as markdown file
    content = lesson["content"]
    filepath = save_markdown_lesson(
        content=content,
        topic=topic["topic"],
        run_id=run_id,
        index=index
    )

    print(f"Saved lesson to {filepath}")

    return lesson

def generate_lessons(run_id: str) -> List[Dict[str, Any]]:
    """
    Generate lessons for all extracted topics.

    Args:
        run_id: Run identifier

    Returns:
        List[Dict]: List of generated lessons
    """
    # Load the topics
    topics_data = load_topics(run_id)
    if not topics_data or "root" not in topics_data:
        raise ValueError(f"No topics found for run ID: {run_id}")

    topics = topics_data["root"]
    print(f"Generating lessons for {len(topics)} topics")

    lessons = []
    for i, topic in enumerate(topics):
        lesson = generate_lesson_for_topic(
            topic=topic,
            run_id=run_id,
            index=i + 1  # 1-based indexing for filenames
        )
        lessons.append(lesson)

    # Save the collection of all lessons
    save_output(
        data={"lessons": lessons},
        pipeline="lessons",
        step="all_lessons",
        run_id=run_id
    )

    return lessons