"""
Build roadmap step that runs expansion and splits the hierarchy.

This module provides a single high-level step `build_roadmap` which coordinates:
  1. expand_and_create_hierarchy (LLM-driven expansion + hierarchy creation)
  2. split_hierarchy (deterministic split into tracks/paths/steps + saves)

It re-uses the existing implementations to avoid duplicating logic while providing
a single callable step for pipeline orchestration.
"""
from typing import Dict, Any, Optional
import traceback

from chains.can_do_steps.expand_and_create_hierarchy import (
    expand_and_create_hierarchy,
    load_expanded_hierarchy,
    get_hierarchy_statistics,
)
from chains.can_do_steps.split_hierarchy import (
    split_hierarchy,
    validate_hierarchy_consistency,
)


def build_roadmap(
    run_id: str,
    force_text: bool = False,
    topic: Optional[str] = None,
    audience: Optional[str] = None,
    purpose: Optional[str] = None,
    style: Optional[str] = None,
    notes: Optional[str] = None,
    language: Optional[str] = None,
    prompt_id: Optional[str] = None,
    input_prefix: Optional[str] = None,
) -> Dict[str, Any]:
    """
    High level step that expands can-do statements into a full hierarchy and then splits it into
    tracks/paths/steps JSON outputs.

    Args:
        run_id: Run identifier
        force_text: If True, reuse raw text output instead of calling LLM
        topic: Optional topic to guide the expansion
        audience: (optional) Intended audience for the statements
        purpose: (optional) Purpose / learning objective focus
        style: (optional) Language style to use for statements
        notes: (optional) Additional guidance or context to include in the prompt
        language: (optional) Output language to request from the model
        prompt_id: (optional) Override identifier for selecting a specific prompt file
        input_prefix: (optional) Input prefix/path for run-scoped input files

    Returns:
        Dict containing merged results: hierarchy, split_summary and statistics

    Note:
        This function will automatically look for input files:
        - input/cando-{run_id}.txt for required can-do statements
        - input/tracks-{run_id}.txt for required tracks
    """
    try:
        print(f"Starting build_roadmap for run_id={run_id}")

        # Step 1: expand and create hierarchy (may call LLM)
        hierarchy = expand_and_create_hierarchy(
            run_id=run_id,
            force_text=force_text,
            topic=topic,
            audience=audience,
            purpose=purpose,
            style=style,
            notes=notes,
            language=language,
            prompt_id=prompt_id,
            input_prefix=input_prefix,
        )

        if not hierarchy:
            raise ValueError("expand_and_create_hierarchy returned no data")

        print("Expanded hierarchy created and saved.")

        # Step 2: split hierarchy into tracks/paths/steps (deterministic)
        split_summary = split_hierarchy(run_id=run_id)

        print("Hierarchy split into tracks/paths/steps.")

        # Step 3: gather statistics and validate (best-effort)
        stats = {}
        try:
            stats = get_hierarchy_statistics(run_id=run_id)
        except Exception:
            # Best-effort fallback: try to load expanded hierarchy and compute minimal stats
            try:
                loaded = load_expanded_hierarchy(run_id=run_id)
                stats = {"tracks": len(loaded.get("tracks", []))} if loaded else {}
            except Exception:
                stats = {}

        # Validate split consistency
        try:
            validation = validate_hierarchy_consistency(run_id=run_id)
        except Exception as e:
            validation = {"error": str(e)}

        result = {
            "hierarchy": hierarchy,
            "split_summary": split_summary,
            "statistics": stats,
            "validation": validation,
        }

        print("✅ build_roadmap completed successfully.")
        return result

    except Exception as e:
        tb = traceback.format_exc()
        print(f"Error in build_roadmap: {e}")
        print(tb)
        raise


if __name__ == "__main__":
    import argparse
    import json

    parser = argparse.ArgumentParser(description="Build roadmap: expand and split hierarchy")
    parser.add_argument("--run-id", required=True, help="Run identifier")
    parser.add_argument("--force-text", action="store_true", help="Reuse raw text output")
    parser.add_argument("--topic", help="Optional topic to guide expansion")
    parser.add_argument("--audience", help="Intended audience for the statements")
    parser.add_argument("--purpose", help="Purpose or learning objective focus")
    parser.add_argument("--style", help="Language style to use for statements")
    parser.add_argument("--notes", help="Additional notes or guidance to include in the prompt")
    parser.add_argument("--language", help="Language to request for generated content")
    parser.add_argument("--prompt-id", help="Override prompt identifier for prompt-aware steps")
    parser.add_argument("--input-prefix", default="input", help="Input prefix/path for can-do input files")
    args = parser.parse_args()

    res = build_roadmap(
        run_id=args.run_id,
        force_text=args.force_text,
        topic=args.topic,
        audience=args.audience,
        purpose=args.purpose,
        style=args.style,
        notes=args.notes,
        language=args.language,
        prompt_id=args.prompt_id,
        input_prefix=args.input_prefix,
    )
    summary = {
        "result_summary": {
            "tracks": res.get("statistics", {}).get("tracks"),
            "split": res.get("split_summary"),
        }
    }
    print(json.dumps(summary, indent=2))