"""
Roadmap tracks component for the Lumabit lesson-generation pipeline.
Tracks are the highest level in the roadmap hierarchy.
"""
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

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

    Args:
        output: Raw output from the LLM

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

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

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

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

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

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

def generate_tracks(run_id: str, force_text: bool = False) -> Dict[str, Any]:
    """
    Generate roadmap tracks.

    Args:
        run_id: Run identifier
        force_text: If True, use existing raw text output instead of calling API

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

        if existing_output:
            print(f"Using existing raw output for tracks in roadmap/{run_id}")
            parsed_tracks = parse_tracks(existing_output)

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

            return parsed_tracks

    print(f"Generating roadmap tracks for run ID: {run_id}")
    result = build_chain(
        chain_name="tracks",
        pipeline="roadmap",
        run_id=run_id,
        input_variables={}  # No additional variables needed
    )

    # Parse the tracks from the result
    parsed_tracks = parse_tracks(result["output"])

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

    return parsed_tracks

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

    Args:
        run_id: Run identifier

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