#!/usr/bin/env python3
"""
MassGen Command Line Interface

A clean CLI for MassGen with file-based configuration support.
Supports both interactive mode and single-question mode.

Usage examples:
    # Use YAML/JSON configuration file
    python -m massgen.cli --config config.yaml "What is the capital of France?"

    # Quick setup with backend and model
    python -m massgen.cli --backend openai --model gpt-4o-mini "What is 2+2?"

    # Interactive mode
    python -m massgen.cli --config config.yaml

    # Multiple agents from config
    python -m massgen.cli --config multi_agent.yaml "Compare different approaches to renewable energy"
"""

import argparse
import asyncio
import json
import os
import sys
import yaml
from pathlib import Path
from typing import Dict, Any, Optional, List

from .utils import MODEL_MAPPINGS, get_backend_type_from_model


# Load environment variables from .env file
def load_env_file():
    """Load environment variables from .env file if it exists."""
    env_file = Path(".env")
    if env_file.exists():
        with open(env_file, "r") as f:
            for line in f:
                line = line.strip()
                if line and not line.startswith("#") and "=" in line:
                    key, value = line.split("=", 1)
                    # Remove quotes if present
                    value = value.strip("\"'")
                    os.environ[key] = value


# Load .env file at module import
load_env_file()

# Add project root to path for imports
project_root = Path(__file__).parent.parent.parent.parent
sys.path.insert(0, str(project_root))

from massgen.backend.response import ResponseBackend
from massgen.backend.grok import GrokBackend
from massgen.backend.claude import ClaudeBackend
from massgen.backend.gemini import GeminiBackend
from massgen.chat_agent import SingleAgent, ConfigurableAgent
from massgen.agent_config import AgentConfig
from massgen.orchestrator import Orchestrator
from massgen.frontend.coordination_ui import CoordinationUI

# Color constants for terminal output
BRIGHT_CYAN = "\033[96m"
BRIGHT_BLUE = "\033[94m"
BRIGHT_GREEN = "\033[92m"
BRIGHT_YELLOW = "\033[93m"
BRIGHT_MAGENTA = "\033[95m"
BRIGHT_RED = "\033[91m"
BRIGHT_WHITE = "\033[97m"
RESET = "\033[0m"
BOLD = "\033[1m"


class ConfigurationError(Exception):
    """Configuration error for CLI."""

    pass


def load_config_file(config_path: str) -> Dict[str, Any]:
    """Load configuration from YAML or JSON file."""
    path = Path(config_path)

    # If file doesn't exist in current path, try massgen/configs/ directory
    if not path.exists():
        # Try in massgen/configs/ directory
        configs_path = Path(__file__).parent / "configs" / path.name
        if configs_path.exists():
            path = configs_path
        else:
            raise ConfigurationError(
                f"Configuration file not found: {config_path} (also checked {configs_path})"
            )

    try:
        with open(path, "r", encoding="utf-8") as f:
            if path.suffix.lower() in [".yaml", ".yml"]:
                return yaml.safe_load(f)
            elif path.suffix.lower() == ".json":
                return json.load(f)
            else:
                raise ConfigurationError(
                    f"Unsupported config file format: {path.suffix}"
                )
    except Exception as e:
        raise ConfigurationError(f"Error reading config file: {e}")


def create_backend(backend_type: str, **kwargs) -> Any:
    """Create backend instance from type and parameters."""
    backend_type = backend_type.lower()

    if backend_type == "openai":
        api_key = kwargs.get("api_key") or os.getenv("OPENAI_API_KEY")
        if not api_key:
            raise ConfigurationError(
                "OpenAI API key not found. Set OPENAI_API_KEY or provide in config."
            )
        return ResponseBackend(api_key=api_key)

    elif backend_type == "grok":
        api_key = kwargs.get("api_key") or os.getenv("XAI_API_KEY")
        if not api_key:
            raise ConfigurationError(
                "Grok API key not found. Set XAI_API_KEY or provide in config."
            )
        return GrokBackend(api_key=api_key)

    elif backend_type == "claude":
        api_key = kwargs.get("api_key") or os.getenv("ANTHROPIC_API_KEY")
        if not api_key:
            raise ConfigurationError(
                "Claude API key not found. Set ANTHROPIC_API_KEY or provide in config."
            )
        return ClaudeBackend(api_key=api_key)

    elif backend_type == "gemini":
        api_key = (
            kwargs.get("api_key")
            or os.getenv("GOOGLE_API_KEY")
            or os.getenv("GEMINI_API_KEY")
        )
        if not api_key:
            raise ConfigurationError(
                "Gemini API key not found. Set GOOGLE_API_KEY or provide in config."
            )
        return GeminiBackend(api_key=api_key)

    else:
        raise ConfigurationError(f"Unsupported backend type: {backend_type}")


def create_agents_from_config(config: Dict[str, Any]) -> Dict[str, ConfigurableAgent]:
    """Create agents from configuration."""
    agents = {}

    # Handle single agent configuration
    if "agent" in config:
        agent_config_data = config["agent"]
        backend_config = agent_config_data.get("backend", {})

        # Infer backend type from model if not explicitly provided
        if "type" not in backend_config and "model" in backend_config:
            backend_type = get_backend_type_from_model(backend_config["model"])
        else:
            backend_type = backend_config.get("type")
            if not backend_type:
                raise ConfigurationError(
                    "Backend type must be specified or inferrable from model"
                )

        backend = create_backend(backend_type, **backend_config)

        # Create proper AgentConfig with backend_params
        if backend_type.lower() == "openai":
            agent_config = AgentConfig.create_openai_config(
                **{k: v for k, v in backend_config.items() if k != "type"}
            )
        elif backend_type.lower() == "claude":
            agent_config = AgentConfig.create_claude_config(
                **{k: v for k, v in backend_config.items() if k != "type"}
            )
        elif backend_type.lower() == "grok":
            agent_config = AgentConfig.create_grok_config(
                **{k: v for k, v in backend_config.items() if k != "type"}
            )
        elif backend_type.lower() == "gemini":
            agent_config = AgentConfig.create_gemini_config(
                **{k: v for k, v in backend_config.items() if k != "type"}
            )
        else:
            # Fallback to basic config
            agent_config = AgentConfig(backend_params=backend_config)

        # Set agent ID and system message
        agent_config.agent_id = agent_config_data.get("id", "agent1")
        agent_config.custom_system_instruction = agent_config_data.get("system_message")

        agent = ConfigurableAgent(config=agent_config, backend=backend)
        agents[agent.agent_id] = agent

    # Handle multiple agents configuration
    elif "agents" in config:
        for agent_config_data in config["agents"]:
            backend_config = agent_config_data.get("backend", {})

            # Infer backend type from model if not explicitly provided
            if "type" not in backend_config and "model" in backend_config:
                backend_type = get_backend_type_from_model(backend_config["model"])
            else:
                backend_type = backend_config.get("type")
                if not backend_type:
                    raise ConfigurationError(
                        "Backend type must be specified or inferrable from model"
                    )

            backend = create_backend(backend_type, **backend_config)

            # Create proper AgentConfig with backend_params
            if backend_type.lower() == "openai":
                agent_config = AgentConfig.create_openai_config(
                    **{k: v for k, v in backend_config.items() if k != "type"}
                )
            elif backend_type.lower() == "claude":
                agent_config = AgentConfig.create_claude_config(
                    **{k: v for k, v in backend_config.items() if k != "type"}
                )
            elif backend_type.lower() == "grok":
                agent_config = AgentConfig.create_grok_config(
                    **{k: v for k, v in backend_config.items() if k != "type"}
                )
            else:
                # Fallback to basic config
                agent_config = AgentConfig(backend_params=backend_config)

            # Set agent ID and system message
            agent_config.agent_id = agent_config_data.get("id", f"agent{len(agents)+1}")
            agent_config.custom_system_instruction = agent_config_data.get(
                "system_message"
            )

            agent = ConfigurableAgent(config=agent_config, backend=backend)
            agents[agent.agent_id] = agent

    else:
        raise ConfigurationError(
            "Configuration must contain either 'agent' or 'agents' section"
        )

    return agents


def create_simple_config(
    backend_type: str, model: str, system_message: Optional[str] = None
) -> Dict[str, Any]:
    """Create a simple single-agent configuration."""
    return {
        "agent": {
            "id": "agent1",
            "backend": {"type": backend_type, "model": model},
            "system_message": system_message or "You are a helpful AI assistant.",
        },
        "ui": {"display_type": "rich_terminal", "logging_enabled": True},
    }


async def run_question_with_history(
    question: str,
    agents: Dict[str, SingleAgent],
    ui_config: Dict[str, Any],
    history: List[Dict[str, Any]],
) -> str:
    """Run MassGen with a question and conversation history."""
    # Build messages including history
    messages = history.copy()
    messages.append({"role": "user", "content": question})

    # Check if we should use orchestrator for single agents (default: False for backward compatibility)
    use_orchestrator_for_single = ui_config.get(
        "use_orchestrator_for_single_agent", True
    )

    if len(agents) == 1 and not use_orchestrator_for_single:
        # Single agent mode with history
        agent = next(iter(agents.values()))
        print(f"\n🤖 {BRIGHT_CYAN}Single Agent Mode{RESET}", flush=True)
        print(f"Agent: {agent.agent_id}", flush=True)
        if history:
            print(f"History: {len(history)//2} previous exchanges", flush=True)
        print(f"Question: {BRIGHT_WHITE}{question}{RESET}", flush=True)
        print("\n" + "=" * 60, flush=True)

        response_content = ""

        async for chunk in agent.chat(messages):
            if chunk.type == "content" and chunk.content:
                response_content += chunk.content
                print(chunk.content, end="", flush=True)
            elif chunk.type == "builtin_tool_results":
                # Skip builtin_tool_results to avoid duplication with real-time streaming
                # The backends already show tool status during execution
                continue
            elif chunk.type == "error":
                print(f"\n❌ Error: {chunk.error}", flush=True)
                return ""

        print("\n" + "=" * 60, flush=True)
        return response_content

    else:
        # Multi-agent mode with history
        orchestrator = Orchestrator(agents=agents)
        # Create a fresh UI instance for each question to ensure clean state
        ui = CoordinationUI(
            display_type=ui_config.get("display_type", "rich_terminal"),
            logging_enabled=ui_config.get("logging_enabled", True),
        )

        print(f"\n🤖 {BRIGHT_CYAN}Multi-Agent Mode{RESET}", flush=True)
        print(f"Agents: {', '.join(agents.keys())}", flush=True)
        if history:
            print(f"History: {len(history)//2} previous exchanges", flush=True)
        print(f"Question: {BRIGHT_WHITE}{question}{RESET}", flush=True)
        print("\n" + "=" * 60, flush=True)

        # For multi-agent with history, we need to use a different approach
        # that maintains coordination UI display while supporting conversation context

        if history and len(history) > 0:
            # Use coordination UI with conversation context
            # Extract current question from messages
            current_question = (
                messages[-1].get("content", question) if messages else question
            )

            # Pass the full message context to the UI coordination
            response_content = await ui.coordinate_with_context(
                orchestrator, current_question, messages
            )
        else:
            # Standard coordination for new conversations
            response_content = await ui.coordinate(orchestrator, question)

        return response_content


async def run_single_question(
    question: str, agents: Dict[str, SingleAgent], ui_config: Dict[str, Any]
) -> str:
    """Run MassGen with a single question."""
    # Check if we should use orchestrator for single agents (default: False for backward compatibility)
    use_orchestrator_for_single = ui_config.get(
        "use_orchestrator_for_single_agent", True
    )

    if len(agents) == 1 and not use_orchestrator_for_single:
        # Single agent mode with existing SimpleDisplay frontend
        agent = next(iter(agents.values()))

        print(f"\n🤖 {BRIGHT_CYAN}Single Agent Mode{RESET}", flush=True)
        print(f"Agent: {agent.agent_id}", flush=True)
        print(f"Question: {BRIGHT_WHITE}{question}{RESET}", flush=True)
        print("\n" + "=" * 60, flush=True)

        messages = [{"role": "user", "content": question}]
        response_content = ""

        async for chunk in agent.chat(messages):
            if chunk.type == "content" and chunk.content:
                response_content += chunk.content
                print(chunk.content, end="", flush=True)
            elif chunk.type == "builtin_tool_results":
                # Skip builtin_tool_results to avoid duplication with real-time streaming
                continue
            elif chunk.type == "error":
                print(f"\n❌ Error: {chunk.error}", flush=True)
                return ""

        print("\n" + "=" * 60, flush=True)
        return response_content

    else:
        # Multi-agent mode
        orchestrator = Orchestrator(agents=agents)
        # Create a fresh UI instance for each question to ensure clean state
        ui = CoordinationUI(
            display_type=ui_config.get("display_type", "rich_terminal"),
            logging_enabled=ui_config.get("logging_enabled", True),
        )

        print(f"\n🤖 {BRIGHT_CYAN}Multi-Agent Mode{RESET}", flush=True)
        print(f"Agents: {', '.join(agents.keys())}", flush=True)
        print(f"Question: {BRIGHT_WHITE}{question}{RESET}", flush=True)
        print("\n" + "=" * 60, flush=True)

        final_response = await ui.coordinate(orchestrator, question)
        return final_response


def print_help_messages():
    print(
        "\n💬 Type your questions below. Use slash commands or press Ctrl+C to stop.",
        flush=True,
    )
    print("💡 Commands: /quit, /exit, /reset, /help", flush=True)
    print("=" * 60, flush=True)


async def run_interactive_mode(
    agents: Dict[str, SingleAgent], ui_config: Dict[str, Any]
):
    """Run MassGen in interactive mode with conversation history."""
    print(f"\n{BRIGHT_CYAN}🤖 MassGen Interactive Mode{RESET}", flush=True)
    print("="*60, flush=True)
    
    # Display configuration
    print(f"📋 {BRIGHT_YELLOW}Configuration:{RESET}", flush=True)
    print(f"   Agents: {len(agents)}", flush=True)
    for agent_id, agent in agents.items():
        backend_name = agent.backend.__class__.__name__.replace("Backend", "")
        print(f"   • {agent_id}: {backend_name}", flush=True)

    use_orchestrator_for_single = ui_config.get(
        "use_orchestrator_for_single_agent", True
    )
    if len(agents) == 1:
        mode = (
            "Single Agent (Orchestrator)"
            if use_orchestrator_for_single
            else "Single Agent (Direct)"
        )
    else:
        mode = "Multi-Agent Coordination"
    print(f"   Mode: {mode}", flush=True)
    print(f"   UI: {ui_config.get('display_type', 'rich_terminal')}", flush=True)

    print_help_messages()

    # Maintain conversation history
    conversation_history = []

    try:
        while True:
            try:
                question = input(f"\n{BRIGHT_BLUE}👤 User:{RESET} ").strip()

                # Handle slash commands
                if question.startswith("/"):
                    command = question.lower()

                    if command in ["/quit", "/exit", "/q"]:
                        print("👋 Goodbye!", flush=True)
                        break
                    elif command in ["/reset", "/clear"]:
                        conversation_history = []
                        # Reset all agents
                        for agent in agents.values():
                            agent.reset()
                        print(
                            f"{BRIGHT_YELLOW}🔄 Conversation history cleared!{RESET}",
                            flush=True,
                        )
                        continue
                    elif command in ["/help", "/h"]:
                        print(
                            f"\n{BRIGHT_CYAN}📚 Available Commands:{RESET}", flush=True
                        )
                        print("   /quit, /exit, /q     - Exit the program", flush=True)
                        print(
                            "   /reset, /clear       - Clear conversation history",
                            flush=True,
                        )
                        print(
                            "   /help, /h            - Show this help message",
                            flush=True,
                        )
                        print(
                            "   /status              - Show current status", flush=True
                        )
                        continue
                    elif command == "/status":
                        print(f"\n{BRIGHT_CYAN}📊 Current Status:{RESET}", flush=True)
                        print(
                            f"   Agents: {len(agents)} ({', '.join(agents.keys())})",
                            flush=True,
                        )
                        use_orch_single = ui_config.get(
                            "use_orchestrator_for_single_agent", True
                        )
                        if len(agents) == 1:
                            mode_display = (
                                "Single Agent (Orchestrator)"
                                if use_orch_single
                                else "Single Agent (Direct)"
                            )
                        else:
                            mode_display = "Multi-Agent"
                        print(f"   Mode: {mode_display}", flush=True)
                        print(
                            f"   History: {len(conversation_history)//2} exchanges",
                            flush=True,
                        )
                        continue
                    else:
                        print(f"❓ Unknown command: {command}", flush=True)
                        print("💡 Type /help for available commands", flush=True)
                        continue

                # Handle legacy plain text commands for backwards compatibility
                if question.lower() in ["quit", "exit", "q"]:
                    print("👋 Goodbye!")
                    break

                if question.lower() in ["reset", "clear"]:
                    conversation_history = []
                    for agent in agents.values():
                        agent.reset()
                    print(f"{BRIGHT_YELLOW}🔄 Conversation history cleared!{RESET}")
                    continue

                if not question:
                    print(
                        "Please enter a question or type /help for commands.",
                        flush=True,
                    )
                    continue

                print(f"\n🔄 {BRIGHT_YELLOW}Processing...{RESET}", flush=True)

                response = await run_question_with_history(
                    question, agents, ui_config, conversation_history
                )

                if response:
                    # Add to conversation history
                    conversation_history.append({"role": "user", "content": question})
                    conversation_history.append(
                        {"role": "assistant", "content": response}
                    )
                    print(f"\n{BRIGHT_GREEN}✅ Complete!{RESET}", flush=True)
                    print(
                        f"{BRIGHT_CYAN}💭 History: {len(conversation_history)//2} exchanges{RESET}",
                        flush=True,
                    )
                    print_help_messages()

                else:
                    print(f"\n{BRIGHT_RED}❌ No response generated{RESET}", flush=True)

            except KeyboardInterrupt:
                print("\n👋 Goodbye!")
                break
            except Exception as e:
                print(f"❌ Error: {e}", flush=True)
                print("Please try again or type /quit to exit.", flush=True)

    except KeyboardInterrupt:
        print("\n👋 Goodbye!")

async def main():
    """Main CLI entry point."""
    parser = argparse.ArgumentParser(
        description="MassGen - Multi-Agent Coordination CLI",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Use configuration file
  python -m massgen.cli --config config.yaml "What is machine learning?"
  
  # Quick single agent setup
  python -m massgen.cli --backend openai --model gpt-4o-mini "Explain quantum computing"
  python -m massgen.cli --backend claude --model claude-sonnet-4-20250514 "Analyze this data"
  
  # Interactive mode
  python -m massgen.cli --config config.yaml
  
  # Create sample configurations
  python -m massgen.cli --create-samples

Environment Variables:
  OPENAI_API_KEY      - Required for OpenAI backend
  XAI_API_KEY         - Required for Grok backend  
  ANTHROPIC_API_KEY   - Required for Claude backend
        """,
    )

    # Question (optional for interactive mode)
    parser.add_argument(
        "question",
        nargs="?",
        help="Question to ask (optional - if not provided, enters interactive mode)",
    )

    # Configuration options
    config_group = parser.add_mutually_exclusive_group()
    config_group.add_argument(
        "--config", type=str, help="Path to YAML/JSON configuration file"
    )
    config_group.add_argument(
        "--backend",
        type=str,
        choices=["openai", "grok", "claude", "gemini"],
        help="Backend type for quick setup",
    )

    # Quick setup options
    parser.add_argument(
        "--model",
        type=str,
        default="gpt-4o-mini",
        help="Model name for quick setup (default: gpt-4o-mini)",
    )
    parser.add_argument(
        "--system-message", type=str, help="System message for quick setup"
    )

    # UI options
    parser.add_argument(
        "--no-display", action="store_true", help="Disable visual coordination display"
    )
    parser.add_argument("--no-logs", action="store_true", help="Disable logging")

    args = parser.parse_args()

    # Validate arguments
    if not args.backend:
        if not args.model and not args.config:
            parser.error(
                "If there is not --backend, either --config or --model must be specified"
            )

    try:
        # Load or create configuration
        if args.config:
            config = load_config_file(args.config)
        else:
            model = args.model
            if args.backend:
                backend = args.backend
            else:
                backend = get_backend_type_from_model(model=model)
            if args.system_message:
                system_message = args.system_message
            else:
                system_message = None
            config = create_simple_config(
                backend_type=backend, model=model, system_message=system_message
            )

        # Apply command-line overrides
        ui_config = config.get("ui", {})
        if args.no_display:
            ui_config["display_type"] = "simple"
        if args.no_logs:
            ui_config["logging_enabled"] = False

        # Create agents
        agents = create_agents_from_config(config)

        if not agents:
            raise ConfigurationError("No agents configured")

        # Run mode based on whether question was provided
        if args.question:
            response = await run_single_question(args.question, agents, ui_config)
            # if response:
            #     print(f"\n{BRIGHT_GREEN}Final Response:{RESET}", flush=True)
            #     print(f"{response}", flush=True)
        else:
            await run_interactive_mode(agents, ui_config)

    except ConfigurationError as e:
        print(f"❌ Configuration error: {e}", flush=True)
        sys.exit(1)
    except KeyboardInterrupt:
        print("\n👋 Goodbye!", flush=True)
    except Exception as e:
        print(f"❌ Error: {e}", flush=True)
        sys.exit(1)


if __name__ == "__main__":
    asyncio.run(main())
