"""
Roadmap bits component for the Lumabit lesson-generation pipeline.
Bits are the smallest units of the roadmap, representing specific learning tasks within steps.
"""
import os
import json
from typing import Dict, Any, List, Optional

from utils.io import save_output, load_latest_output
from chains.base import build_chain, default_json_parser, parse_output
from chains.roadmap.steps import load_steps

def parse_bits(output: str) -> Dict[str, Any]:
    """
    Parse the bits from the LLM output.

    Args:
        output: Raw output from the LLM

    Returns:
        Dict: Parsed bits data
    """
    try:
        # Parse the JSON from the output
        parsed = parse_output(output, default_json_parser)

        # Validate the expected structure
        if "bits" not in parsed:
            raise ValueError("Expected 'bits' key in parsed output")

        bits = parsed["bits"]
        if not isinstance(bits, list):
            raise ValueError("Expected 'bits' to be a list")

        # Validate each bit has the required fields
        required_fields = ["id", "step_id", "title", "content", "order", "type"]
        for i, bit in enumerate(bits):
            if not isinstance(bit, dict):
                raise ValueError(f"Bit at index {i} is not a dictionary")

            for field in required_fields:
                if field not in bit:
                    raise ValueError(f"Bit at index {i} is missing required field '{field}'")

        return parsed
    except Exception as e:
        print(f"Error parsing bits: {e}")
        raise

def generate_bits(
    run_id: str,
    step_id: Optional[str] = None,
    force_text: bool = False
) -> Dict[str, Any]:
    """
    Generate roadmap bits for a specific step or all steps.

    Args:
        run_id: Run identifier
        step_id: Optional step ID to generate bits for
        force_text: If True, use existing raw text output instead of calling API

    Returns:
        Dict: Generated bits data
    """
    # Check if we should use existing raw output
    if force_text:
        existing_output = load_latest_output(
            pipeline="roadmap",
            step="bits",
            run_id=run_id,
            as_text=True,
            raw=True
        )

        if existing_output:
            print(f"Using existing raw output for bits in roadmap/{run_id}")
            parsed_bits = parse_bits(existing_output)

            # Save the parsed output
            save_output(
                data=parsed_bits,
                pipeline="roadmap",
                step="bits",
                run_id=run_id
            )

            return parsed_bits

    # Load the steps if we need them
    steps_data = load_steps(run_id)
    if not steps_data:
        raise ValueError(f"No steps found for run ID: {run_id}")

    # Prepare input variables
    input_variables = {}

    # If a specific step ID is provided, filter the steps
    if step_id:
        steps = [s for s in steps_data["steps"] if s["id"] == step_id]
        if not steps:
            raise ValueError(f"Step with ID {step_id} not found")

        input_variables["step"] = steps[0]
    else:
        input_variables["steps"] = steps_data["steps"]

    print(f"Generating roadmap bits for run ID: {run_id}")
    result = build_chain(
        chain_name="bits",
        pipeline="roadmap",
        run_id=run_id,
        input_variables=input_variables
    )

    # Parse the bits from the result
    parsed_bits = parse_bits(result["output"])

    # Save the parsed output
    save_output(
        data=parsed_bits,
        pipeline="roadmap",
        step="bits",
        run_id=run_id
    )

    return parsed_bits

def load_bits(run_id: str) -> Optional[Dict[str, Any]]:
    """
    Load previously generated bits.

    Args:
        run_id: Run identifier

    Returns:
        Dict: Previously generated bits, or None if not found
    """
    return load_latest_output(
        pipeline="roadmap",
        step="bits",
        run_id=run_id
    )

def generate_complete_roadmap(run_id: str) -> Dict[str, Any]:
    """
    Generate a complete roadmap with all components.

    Args:
        run_id: Run identifier

    Returns:
        Dict: Complete roadmap data
    """
    from chains.roadmap.tracks import generate_tracks
    from chains.roadmap.paths import generate_paths
    from chains.roadmap.steps import generate_steps

    print(f"Generating complete roadmap for run ID: {run_id}")

    # Generate tracks
    tracks_data = generate_tracks(run_id)

    # Generate paths for all tracks
    paths_data = generate_paths(run_id)

    # Generate steps for all paths
    steps_data = generate_steps(run_id)

    # Generate bits for all steps
    bits_data = generate_bits(run_id)

    # Combine all data into a complete roadmap
    roadmap = {
        "tracks": tracks_data["tracks"],
        "paths": paths_data["paths"],
        "steps": steps_data["steps"],
        "bits": bits_data["bits"]
    }

    # Save the complete roadmap
    save_output(
        data=roadmap,
        pipeline="roadmap",
        step="complete",
        run_id=run_id
    )

    return roadmap