"""
Roadmap steps component for the Lumabit lesson-generation pipeline.
Steps are specific learning objectives within a path.
"""
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.paths import load_paths

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

    Args:
        output: Raw output from the LLM

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

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

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

        # Validate each step has the required fields
        required_fields = ["id", "path_id", "title", "description", "order"]
        for i, step in enumerate(steps):
            if not isinstance(step, dict):
                raise ValueError(f"Step at index {i} is not a dictionary")

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

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

def generate_steps(
    run_id: str,
    path_id: Optional[str] = None,
    force_text: bool = False
) -> Dict[str, Any]:
    """
    Generate roadmap steps for a specific path or all paths.

    Args:
        run_id: Run identifier
        path_id: Optional path ID to generate steps for
        force_text: If True, use existing raw text output instead of calling API

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

        if existing_output:
            print(f"Using existing raw output for steps in roadmap/{run_id}")
            parsed_steps = parse_steps(existing_output)

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

            return parsed_steps

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

    # Prepare input variables
    input_variables = {}

    # If a specific path ID is provided, filter the paths
    if path_id:
        paths = [p for p in paths_data["paths"] if p["id"] == path_id]
        if not paths:
            raise ValueError(f"Path with ID {path_id} not found")

        input_variables["path"] = paths[0]
    else:
        input_variables["paths"] = paths_data["paths"]

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

    # Parse the steps from the result
    parsed_steps = parse_steps(result["output"])

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

    return parsed_steps

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

    Args:
        run_id: Run identifier

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