import argparse
import os
import json
import sys
import time
from story_builder import Engine, StoryGraph, EngineState
from story_builder.utils.loaders import load_graph

def visualize_command(args):
    """Command for story graph visualization."""
    from story_builder.export.render import write_dot, render_dot_to_image, open_file

    if not os.path.exists(args.path):
        print(f"Error: Path {args.path} does not exist.")
        sys.exit(2)

    try:
        with open(args.path, 'r', encoding='utf-8-sig') as f:
            data = json.load(f)
    except Exception as e:
        print(f"Error: Failed to parse JSON file (UTF-8 BOM supported): {e}")
        sys.exit(2)

    try:
        graph = StoryGraph(**data)
        
        # Paths logic
        base, _ = os.path.splitext(args.path)
        out_dot = f"{base}.dot"
        out_img = args.out or f"{base}.{args.format}"
        
        # 1. Output DOT
        write_dot(graph, out_dot)
        
        if getattr(args, 'dot_only', False):
            return
    except Exception as e:
        print(f"Error during DOT export: {e}")
        sys.exit(2)

    try:
        # 2. Render to image
        success = render_dot_to_image(out_dot, out_img, format=args.format, engine=args.engine)
        
        if success:
            print(f"Successfully visualized to: {out_img}")
            if not getattr(args, 'no_open', False):
                open_file(out_img)
        else:
            print("\nHint: Only the .dot file was generated. To render it as an image, please:")
            print("  - pip install 'story-builder[viz]'")
            print("  - Ensure GraphViz system binaries are installed (http://graphviz.org)")
    except Exception as e:
        print(f"Error during rendering: {e}")
        # sys.exit(2) # Don't exit on render fail if DOT was ok

def play_command(args):
    """Command for playing through a story."""
    if not args.path:
        print("Error: path is required for play.")
        sys.exit(2)
        
    if not os.path.exists(args.path):
        print(f"Error: Story file '{args.path}' not found.")
        sys.exit(2)
        
    try:
        story_graph = load_graph(args.path)
    except Exception as e:
        print(f"Error: {e}")
        sys.exit(2)
        
    engine = Engine(story_graph)
    
    # Start game
    state = EngineState(current_node_id=story_graph.start_node_id, variables={})
    
    print(f"\nStarting Adventure: {story_graph.metadata.get('title', 'Unknown')}")
    print("-" * 50)
    
    while True:
        node, triggers, _ = engine.run(state.current_node_id, state.variables)
        
        print(f"\nNode: {node.title}")
        print(f"Body: {node.body}")
        
        if not triggers:
            print("\n[GAME OVER] No choices left.")
            break
            
        print("\nAvailable Actions:")
        for i, t in enumerate(triggers, start=1):
            print(f"  {i}. {t}")
        
        try:
            choice = input("\n> ")
            if choice.lower() in ['exit', 'quit']:
                break
            
            # Check for numerical choice
            if choice.isdigit() and 0 < int(choice) <= len(triggers):
                selected_trigger = triggers[int(choice)-1]
            else:
                selected_trigger = choice
            
            # Process trigger
            engine.process_trigger(selected_trigger, state)
            
        except Exception as e:
            print(f"\nError: {e}")
            print("Please try again.")

def init_command(args):
    """Command to generate a starter story template."""
    starter_story = {
        "metadata": {
            "title": "My New Story"
        },
        "nodes": {
            "start": {
                "id": "start",
                "title": "The Beginning",
                "body": "You stand at a crossroads. Where will you go?",
                "choices": [
                    {
                        "trigger": "Go North",
                        "target_node_id": "north_node"
                    },
                    {
                        "trigger": "Go South",
                        "target_node_id": "south_node"
                    }
                ]
            },
            "north_node": {
                "id": "north_node",
                "title": "The North Pole",
                "body": "It is freezing here. You see a polar bear.",
                "choices": []
            },
            "south_node": {
                "id": "south_node",
                "title": "The Sunny South",
                "body": "Beach vibes everywhere. You feel relaxed.",
                "choices": []
            }
        },
        "start_node_id": "start"
    }
    
    filename = "story.json"
    try:
        if os.path.exists(filename) and not getattr(args, 'force', False):
            print(f"Error: {filename} already exists. Use --force to overwrite.")
            sys.exit(1)
            
        with open(filename, 'w', encoding='utf-8') as f:
            json.dump(starter_story, f, indent=2)
        print(f"Created starter story: {filename}")
    except Exception as e:
        print(f"Error creating starter story: {e}")
        sys.exit(1)

def validate_command(args):
    """Command to validate a story graph JSON file."""
    if not os.path.exists(args.path):
        print(f"Error: File {args.path} not found.")
        sys.exit(1)
    
    try:
        graph = load_graph(args.path)
        graph.validate_graph()
        print("Graph valid.")
    except Exception as e:
        print(f"Graph validation error: {e}")
        sys.exit(1)

def export_command(args):
    """Command to export a story graph to other formats."""
    if not os.path.exists(args.path):
        print(f"Error: File {args.path} not found.")
        sys.exit(1)

    try:
        graph = load_graph(args.path)
        base, _ = os.path.splitext(args.path)
        
        if args.format == "ink":
            from story_builder.export.ink_exporter import InkExporter
            out_file = f"{base}.ink"
            InkExporter.export(graph, out_file)
            print(f"Exported to {out_file} (Ink format)")
        else:
            print(f"Format {args.format} not implemented or same as source.")
            
    except Exception as e:
        print(f"Export failed: {e}")
        sys.exit(1)

def generate_command(args):
    """Command to generate a whole story from a prompt using AI."""
    import random
    import io
    from story_builder.export.render import write_dot
    
    print(f"=== Story Builder AI Generator ===")
    
    api_key = args.api_key or os.getenv("OPENROUTER_API_KEY")
    is_mock = getattr(args, 'mock', False) or not api_key

    if is_mock:
        print("Running in MOCK mode (no API call).")

    def build_mock_graph():
        rng = random.Random(args.seed)
        prompt_lower = args.prompt.lower()
        
        verbs = ["investigate", "talk", "sneak", "leave", "hack", "run", "look", "go", "climb", "hide"]
        
        # Simple heuristics for theme
        if any(w in prompt_lower for w in ["cyberpunk", "neon", "sci-fi", "future"]):
            title_prefix = "Sector"
            body_templates = ["Neon lights flicker.", "Rain-slicked streets reflect digital signs.", "A drone hums overhead."]
        elif any(w in prompt_lower for w in ["fantasy", "magic", "dragon", "sword", "quest"]):
            title_prefix = "Realm"
            body_templates = ["The air smells of ancient magic.", "Stone ruins tower over the path.", "The wind whispers through the oak trees."]
        else:
            title_prefix = "Area"
            body_templates = ["It is quiet here.", "The shadow grows longer.", "Dust motes dance in the light."]
            
        nodes = {
            "start": {
                "id": "start",
                "title": f"Start: {title_prefix} 0",
                "body": rng.choice(body_templates),
                "choices": []
            }
        }
        
        target_nodes = max(3, args.nodes)
        queue = [("start", 0)]
        current_idx = 1
        
        while queue and len(nodes) < target_nodes:
            parent_id, depth = queue.pop(0)
            
            if depth >= args.depth:
                continue
                
            # Try to aim for branching factor K
            num_choices = rng.randint(1, min(args.branching + 1, target_nodes - len(nodes) + 1))
            
            # Limit branching if we are close to the node limit
            remaining = target_nodes - len(nodes)
            if remaining <= 0:
                break
            
            num_choices = min(num_choices, remaining)
            
            # Ensure we don't pick more verbs than we have
            selected_verbs = rng.sample(verbs, min(num_choices, len(verbs)))
                
            for verb in selected_verbs:
                child_id = f"node_{current_idx}"
                current_idx += 1
                
                nodes[parent_id]["choices"].append({
                    "trigger": verb,
                    "target_node_id": child_id
                })
                
                nodes[child_id] = {
                    "id": child_id,
                    "title": f"{title_prefix} {current_idx}",
                    "body": f"{rng.choice(body_templates)} ({current_idx})",
                    "choices": []
                }
                queue.append((child_id, depth + 1))
                
        # Enforce minimum 3 nodes invariant
        while len(nodes) < 3:
            parent_id = "start"
            child_id = f"node_{current_idx}"
            current_idx += 1
            
            existing_start_triggers = [c.get("trigger") for c in nodes["start"]["choices"]]
            available_verbs = [v for v in verbs if v not in existing_start_triggers]
            verb = rng.choice(available_verbs) if available_verbs else f"verb_{current_idx}"
            
            nodes[parent_id]["choices"].append({"trigger": verb, "target_node_id": child_id})
            nodes[child_id] = {
                "id": child_id,
                "title": f"{title_prefix} {current_idx}",
                "body": f"{rng.choice(body_templates)} ({current_idx})",
                "choices": []
            }

        graph_data = {
            "metadata": {"title": f"MOCK: {args.prompt[:20]}..."},
            "nodes": nodes,
            "start_node_id": "start"
        }
        return StoryGraph(**graph_data)

    try:
        print("[1/4] Generating story graph")
        time.sleep(0.3)
        
        graph = None
        if is_mock:
            graph = build_mock_graph()
        else:
            try:
                from story_builder_pro.ai.provider_openrouter import OpenRouterProvider
                provider = OpenRouterProvider(api_key=api_key)
                
                system_prompt = "Generate a branching story graph in JSON. Return JSON only."
                base_user_prompt = f"Generate a story with about {args.nodes} nodes, maximum depth {args.depth}, and average branching factor {args.branching} about: {args.prompt}. Use format: metadata, nodes (key is id, value has id, title, body, choices: [trigger, target_node_id]), start_node_id."
                user_prompt = base_user_prompt
                
                max_retries = 3
                for attempt in range(max_retries):
                    try:
                        ai_data = provider.generate_json(user_prompt, system_prompt=system_prompt)
                        # Basic fix: Ensure 'to' vs 'target_node_id' if needed, though StoryGraph pydantic alias handles it
                        temp_graph = StoryGraph(**ai_data)
                        temp_graph.validate_graph()
                        if len(temp_graph.nodes) < 3:
                             raise ValueError("Graph has fewer than 3 nodes.")
                        graph = temp_graph
                        break
                    except Exception as e:
                        if attempt < max_retries - 1:
                            user_prompt = base_user_prompt + f"\n\nThe previous attempt failed with error: {e}. Fix the invalid references (all choice 'target_node_id' must exist in 'nodes'), ensure at least 3 nodes, and return valid JSON only."
                
                if not graph:
                    print("AI generation failed; falling back to MOCK generator.")
                    graph = build_mock_graph()
            except Exception as e:
                print(f"AI generation failed ({e}); falling back to MOCK generator.")
                graph = build_mock_graph()

        print("[2/4] Validating graph")
        time.sleep(0.3)
        graph.validate_graph()

        print(f"[3/4] Saving story: {args.output}")
        time.sleep(0.3)
        with open(args.output, "w", encoding="utf-8") as f:
            f.write(graph.model_dump_json(indent=2))
        
        base, _ = os.path.splitext(args.output)
        out_dot = f"{base}.dot"
        print(f"[4/4] Exporting DOT: {out_dot}")
        
        # Suppress noise from internal write_dot if any
        old_stdout = sys.stdout
        sys.stdout = io.StringIO()
        try:
            write_dot(graph, out_dot)
        finally:
            sys.stdout = old_stdout
        
    except Exception as e:
        print(f"Error during generation: {e}")
        sys.exit(1)

def demo_ai_command(args):
    """Command for AI branch generation demo."""
    print("=== Story Builder AI Demo ===")
    
    try:
        import story_builder_pro
        from story_builder_pro.ai.branch_generator import BranchGenerator
        from story_builder_pro.ai.provider_openrouter import OpenRouterProvider
    except ImportError:
        print("This demo requires story_builder_pro")
        sys.exit(1)

    api_key = args.api_key or os.getenv("OPENROUTER_API_KEY")
    is_mock = getattr(args, 'mock', False) or not api_key

    if is_mock:
        print("Running in MOCK mode (no API call).")

    print(f"\n[1/4] Loading base story")
    time.sleep(0.3)
    base_story_path = "examples/minimal_story.json"
    if not os.path.exists(base_story_path):
        print(f"Error: {base_story_path} not found.")
        sys.exit(1)
        
    try:
        graph = load_graph(base_story_path)
        print("Story loaded")
    except Exception as e:
        print(f"Error loading story: {e}")
        sys.exit(1)

    print(f"\n[2/4] Generating AI branches")
    time.sleep(0.3)
    try:
        node_id = "start" if "start" in graph.nodes else graph.start_node_id

        if is_mock:
            mock_ai_data = {
                "choices": [
                    {"trigger": "explore forest", "to": "forest_path"},
                    {"trigger": "enter village", "to": "village_square"},
                    {"trigger": "follow river", "to": "river_bank"}
                ],
                "new_nodes": [
                    {"id": "forest_path", "title": "The Forest Path", "body": "Distant birds chirp in the canopy.", "choices": []},
                    {"id": "village_square", "title": "Village Square", "body": "A bustling market with smells of fresh bread.", "choices": []},
                    {"id": "river_bank", "title": "The River Bank", "body": "The water flows swiftly toward the sea.", "choices": []}
                ]
            }
            generator = BranchGenerator(None)
            updated_graph = generator._merge_and_validate(graph, node_id, mock_ai_data)
            
            print(f"AI created {len(mock_ai_data['choices'])} choices")
            print(f"AI created {len(mock_ai_data['new_nodes'])} nodes")
        else:
            provider = OpenRouterProvider(api_key=api_key)
            generator = BranchGenerator(provider)
            
            nodes_before = len(graph.nodes)
            choices_before = len(graph.nodes[node_id].choices)
            
            updated_graph = generator.generate(
                graph=graph,
                node_id=node_id,
                goal="Add a mysterious branching path in the forest",
                max_choices=3
            )
            
            nodes_after = len(updated_graph.nodes)
            choices_after = len(updated_graph.nodes[node_id].choices)
            
            print(f"AI created {choices_after - choices_before} choices")
            print(f"AI created {nodes_after - nodes_before} nodes")
            
    except Exception as e:
        print(f"Error during AI generation: {e}")
        sys.exit(1)

    print(f"\n[3/4] Saving updated story")
    time.sleep(0.3)
    output_path = "examples/generated_story.json"
    try:
        with open(output_path, 'w', encoding='utf-8') as f:
            f.write(updated_graph.model_dump_json(indent=2))
        print(f"Saved to {output_path}")
    except Exception as e:
        print(f"Error saving story: {e}")
        sys.exit(1)

    print(f"\n[4/4] Visualizing graph")
    time.sleep(0.3)
    try:
        viz_args = argparse.Namespace(path=output_path, dot_only=True, out=None, format="png", no_open=True, engine="dot")
        
        base, _ = os.path.splitext(output_path)
        out_dot = f"{base}.dot"
        
        with open(os.devnull, 'w') as f:
            old_stdout = sys.stdout
            sys.stdout = f
            try:
                visualize_command(viz_args)
            finally:
                sys.stdout = old_stdout
                
        print(f"Graph saved to {os.path.basename(out_dot)}")
    except Exception as e:
        print(f"Error during visualization: {e}")

    print(f"\nDemo complete.")

def main():
    parser = argparse.ArgumentParser(description="Story Builder Narrative Engine CLI")
    
    parser.add_argument("--load", type=str, help="Load an example story JSON (deprecated: use 'play' subcommand)")
    parser.add_argument("--play", action="store_true", help="Start a new adventure (deprecated: use 'play' subcommand)")

    subparsers = parser.add_subparsers(dest="subcommand", help="sub-command help")

    # init
    init_parser = subparsers.add_parser("init", help="Generate a starter story template.")
    init_parser.add_argument("--force", action="store_true", help="Overwrite existing story.json if it exists.")

    # generate
    gen_parser = subparsers.add_parser("generate", help="Generate a branching story graph using AI.")
    gen_parser.add_argument("prompt", help="Story theme or prompt.")
    gen_parser.add_argument("--nodes", type=int, default=6, help="Total target nodes (default: 6).")
    gen_parser.add_argument("--branching", type=int, default=2, help="Average branching factor (default: 2).")
    gen_parser.add_argument("--depth", type=int, default=3, help="Maximum depth from start (default: 3).")
    gen_parser.add_argument("--output", type=str, default="story.json", help="Output JSON file (default: story.json).")
    gen_parser.add_argument("--seed", type=int, default=42, help="Deterministic generation seed (default: 42).")
    gen_parser.add_argument("--mock", action="store_true", help="Force mock mode even if API key is available.")
    gen_parser.add_argument("--api-key", type=str, help="OpenRouter API Key (overrides env OPENROUTER_API_KEY).")

    # validate
    val_parser = subparsers.add_parser("validate", help="Validate a story graph JSON file.")
    val_parser.add_argument("path", help="Path to the story JSON file.")

    # export
    exp_parser = subparsers.add_parser("export", help="Export a story graph to other formats.")
    exp_parser.add_argument("path", help="Path to the story JSON file.")
    exp_parser.add_argument("--format", type=str, default="json", choices=["json", "ink"], help="Export format (default: json).")

    # play
    play_parser = subparsers.add_parser("play", help="Play an interactive story.")
    play_parser.add_argument("path", nargs="?", help="Path to the story JSON file.")

    # visualize
    viz_parser = subparsers.add_parser("visualize", help="Generate a graph visualization from story JSON.")
    viz_parser.add_argument("path", help="Path to the story JSON file.")
    viz_parser.add_argument("--out", type=str, help="Output image path.")
    viz_parser.add_argument("--format", type=str, default="png", choices=["png", "svg"], help="Output format (default: png).")
    viz_parser.add_argument("--no-open", action="store_true", help="Do not auto-open the rendered image.")
    viz_parser.add_argument("--dot-only", action="store_true", help="Only generate the DOT file, skip rendering.")
    viz_parser.add_argument("--engine", type=str, default="dot", choices=["dot", "neato"], help="GraphViz layout engine (default: dot).")

    # demo-ai
    demo_ai_parser = subparsers.add_parser("demo-ai", help="Run an AI demo that generates branches.")
    demo_ai_parser.add_argument("--mock", action="store_true", help="Run in mock mode (no API calls).")
    demo_ai_parser.add_argument("--api-key", type=str, help="OpenRouter API Key.")

    args = parser.parse_args()

    if args.load:
        args.path = args.load
        play_command(args)
        return

    if args.subcommand == "play":
        play_command(args)
    elif args.subcommand == "visualize":
        visualize_command(args)
    elif args.subcommand == "demo-ai":
        demo_ai_command(args)
    elif args.subcommand == "init":
        init_command(args)
    elif args.subcommand == "generate":
        generate_command(args)
    elif args.subcommand == "validate":
        validate_command(args)
    elif args.subcommand == "export":
        export_command(args)
    else:
        parser.print_help()
        sys.exit(0)

if __name__ == "__main__":
    main()
