import argparse
import asyncio
import sys
import time
import uvicorn
from app.core.interceptor import InterceptorProxy
from app.core.config import register_server, list_servers, get_server, APP_NAME

def show_splash():
    """
    애플리케이션 시작 시 시각적인 즐거움을 위한 ASCII 아트 스플래시 화면을 출력합니다.
    """
    splash = r"""
        ________________________________________________
       /  ____________________________________________  \
      |  /                                            \  |
      | |    _      _      __  __    _  _             | |
      | |   | |    | |    |  \/  |  | || |   ___      | |
      | |   | |    | |    | |\/| |  | || |_ / __|     | |
      | |   | |___ | |___ | |  | |  |__   _|\__ \     | |
      | |   |_____||_____||_|  |_|     |_|  |___/     | |
      | |                                             | |
      | |         [ MCP LENS DISCOVERY ]              | |
      |  \____________________________________________/  |
       \________________________________________________/
              \________________________________/
    """
    print(splash)
    time.sleep(0.5)

def run_proxy(target_command: list[str], name: str = None, provider: str = "Unknown", model: str = "DirectCommand"):
    """
    Proxy 모드 실행
    """
    cmd = target_command
    if name and not cmd:
        server = get_server(name)
        if server:
            cmd = [server["command"]] + server["args"]
            print(f"Running registered server '{name}': {' '.join(cmd)}")
        else:
            print(f"Error: Server '{name}' not found.", file=sys.stderr)
            print(f"\nRegistered servers:", file=sys.stderr)
            for s in list_servers():
                print(f"  - {s['name']}", file=sys.stderr)
            print(f"\nTip: Register with: llm4s add-server {name} --command <cmd> -- <args>", file=sys.stderr)
            sys.exit(1)

    if not cmd:
        print("Error: Target command or --name is required for proxy mode.", file=sys.stderr)
        print("\nUsage:", file=sys.stderr)
        print("  llm4s proxy -- node /path/to/mcp-server/index.js", file=sys.stderr)
        print("  llm4s proxy --name my-registered-server", file=sys.stderr)
        print("\nTip: Register servers with:", file=sys.stderr)
        print("  llm4s add-server my-server --command node -- /path/to/server.js", file=sys.stderr)
        print("  llm4s list-servers   # View registered servers", file=sys.stderr)
        sys.exit(1)
        
    # Provider 및 모델명 추출 (메타데이터용)
    provider_name = provider
    server_name = model if model != "DirectCommand" else (name or "DirectCommand")
    
    # 만약 프로바이더가 지정되지 않았다면 (Unknown), 이름이나 명령어에 기반한 자동 프로바이더 판단
    if provider_name == "Unknown":
        full_cmd_str = " ".join(cmd).lower()
        if name:
            name_lower = name.lower()
            if "google" in name_lower or "gemini" in name_lower:
                provider_name = "Google"
            elif "anthropic" in name_lower or "claude" in name_lower:
                provider_name = "Anthropic"
            elif "openai" in name_lower or "gpt" in name_lower:
                provider_name = "OpenAI"
        
        if provider_name == "Unknown":
            if "google" in full_cmd_str or "gemini" in full_cmd_str:
                provider_name = "Google"
            elif "anthropic" in full_cmd_str or "claude" in full_cmd_str:
                provider_name = "Anthropic"
            elif "openai" in full_cmd_str or "gpt" in full_cmd_str:
                provider_name = "OpenAI"

    proxy = InterceptorProxy(cmd, provider=provider_name, model=server_name)
    
    try:
        asyncio.run(proxy.start())
    except KeyboardInterrupt:
        print("\nProxy terminated by user.", file=sys.stderr)

def run_api(host: str, port: int):
    """
    FastAPI 서버 모드 실행
    """
    uvicorn.run("app.api.main:app", host=host, port=port, reload=True)

def main():
    show_splash()
    parser = argparse.ArgumentParser(description="LLM4s (MCP-Lens) CLI")
    subparsers = parser.add_subparsers(dest="mode", required=True, help="Run mode")
    
    # 1. Proxy Mode
    proxy_parser = subparsers.add_parser("proxy", help="Run the stdio interceptor proxy")
    proxy_parser.add_argument("--name", help="Registered server name to run")
    proxy_parser.add_argument("-p", "--provider", default="Unknown", help="Provider name for metadata")
    proxy_parser.add_argument("-m", "--model", default="DirectCommand", help="Model name for metadata")
    proxy_parser.add_argument(
        "target", 
        nargs=argparse.REMAINDER, 
        help="The actual MCP server command to run"
    )
    
    # 2. Add Server
    add_parser = subparsers.add_parser("add-server", help="Register a new MCP server")
    add_parser.add_argument("name", help="Server name")
    add_parser.add_argument("--command", required=True, help="Command to run")
    add_parser.add_argument("--args", nargs=argparse.REMAINDER, help="Arguments for the command")

    # 3. List Servers
    subparsers.add_parser("list-servers", help="List all registered servers")

    # 4. API Mode
    api_parser = subparsers.add_parser("api", help="Run the FastAPI backend server")
    api_parser.add_argument("--host", default="127.0.0.1", help="API server host")
    api_parser.add_argument("--port", type=int, default=8000, help="API server port")

    # 5. TUI Mode
    subparsers.add_parser("tui", help="Run the Textual UI Dashboard")

    # 7. Config Mode (Standardized)
    config_parser = subparsers.add_parser("config", help="Standardized configuration management")
    config_parser.add_argument("--add", nargs=2, metavar=("PROVIDER", "KEY"), help="Add API key for provider (openai, claude, gemini)")
    config_parser.add_argument("--list", action="store_true", help="List all registered providers")

    # 8. Doctor Mode
    subparsers.add_parser("doctor", help="Check system health and configuration status")

    # 9. Test Traffic Mode
    test_parser = subparsers.add_parser("test-traffic", help="Run a long-running LLM traffic simulation")
    test_parser.add_argument("--duration", type=int, default=10, help="Duration in seconds")
    test_parser.add_argument("--provider", default="Unknown", help="Provider name for testing")
    test_parser.add_argument("--model", default="DirectCommand", help="Model name for testing")

    args = parser.parse_args()

    if args.mode == "proxy":
        target = args.target
        if target and target[0] == "--":
            target = target[1:]
        run_proxy(target, name=args.name, provider=args.provider, model=args.model)
        
    elif args.mode == "add-server":
        args_list = args.args
        if args_list and args_list[0] == "--":
            args_list = args_list[1:]
        register_server(args.name, args.command, args_list or [])
        print(f"Server '{args.name}' registered successfully.")

    elif args.mode == "list-servers":
        servers = list_servers()
        if not servers:
            print("No servers registered.")
        else:
            print("Registered Servers:")
            for s in servers:
                print(f"- {s['name']}: {s['command']} {' '.join(s['args'])}")

    elif args.mode == "api":
        run_api(args.host, args.port)
        
    elif args.mode == "tui":
        from app.ui.app import run_tui
        run_tui()

    elif args.mode == "config":
        from app.core.config import register_provider, list_providers
        if args.add:
            provider_alias = args.add[0].lower()
            api_key = args.add[1]
            
            mapping = {
                "openai": ("OpenAI", "OPENAI_API_KEY"),
                "claude": ("Anthropic", "ANTHROPIC_API_KEY"),
                "anthropic": ("Anthropic", "ANTHROPIC_API_KEY"),
                "gemini": ("Google", "GOOGLE_API_KEY"),
                "google": ("Google", "GOOGLE_API_KEY")
            }
            
            if provider_alias in mapping:
                real_name, env_key = mapping[provider_alias]
                register_provider(real_name, env_key, api_key=api_key)
                print(f"Successfully registered API key for {real_name}.")
            else:
                # If not an alias, use as is
                register_provider(args.add[0], f"{args.add[0].upper()}_API_KEY", api_key=api_key)
                print(f"Registered API key for custom provider '{args.add[0]}'.")
        
        if args.list or len(sys.argv) == 2:
            providers = list_providers()
            if not providers:
                print("No API keys registered.")
                print("\nQuick setup:")
                print("  llm4s config --add gemini YOUR_KEY")
                print("  llm4s config --add openai YOUR_KEY")
                print("  llm4s config --add claude YOUR_KEY")
                print("\nOr set environment variables:")
                print("  export GOOGLE_API_KEY=...")
                print("  export OPENAI_API_KEY=...")
                print("  export ANTHROPIC_API_KEY=...")
            else:
                print(f"{'PROVIDER':<15} {'STATUS':<10} {'SOURCE':<8} {'ENV_KEY':<25} {'KEY_HINT'}")
                print("-" * 75)
                for p in providers:
                    source = p.get('source', '?')
                    print(f"{p['name']:<15} {p['status']:<10} {source:<8} {p.get('env_key',''):<25} {p['key_hint']}")

    elif args.mode == "doctor":
        from app.core.config import CONFIG_PATH, DB_PATH, list_providers, list_servers
        import os
        
        print(f"[{APP_NAME} DOCTOR - System Health Check]")
        print("=" * 50)
        
        # 0. Runtime
        print(f"● Python: {sys.version.split()[0]}")
        print(f"● Platform: {sys.platform}")
        
        # 1. Paths
        print(f"\n● Config: {CONFIG_PATH} ({'✓ EXISTS' if CONFIG_PATH.exists() else '✗ NOT FOUND'})")
        db_exists = DB_PATH.exists()
        db_size = ""
        if db_exists:
            size_mb = os.path.getsize(DB_PATH) / (1024 * 1024)
            db_size = f" ({size_mb:.1f} MB)"
        print(f"● Database: {DB_PATH} ({'✓ EXISTS' if db_exists else '✗ NOT FOUND'}){db_size}")
        
        # 2. Providers
        providers = list_providers()
        print(f"\n● Providers ({len(providers)}):")
        if not providers:
            print("  (none registered — run: llm4s config --add gemini YOUR_KEY)")
        for p in providers:
            status_icon = "✓" if p['status'] == "OK" else "✗"
            source = p.get('source', '?')
            print(f"  {status_icon} {p['name']:<15} : {p['status']} (via {source})")
            
        # 3. Servers
        servers = list_servers()
        print(f"\n● Registered Servers ({len(servers)}):")
        if not servers:
            print("  (none — run: llm4s add-server my-server --command node -- /path/to/server.js)")
        for s in servers:
            print(f"  - {s['name']:<15} : {s['command']} {' '.join(s['args'])}")
            
        # 4. Environment
        print("\n● Environment Variables:")
        for key in ["OPENAI_API_KEY", "ANTHROPIC_API_KEY", "GOOGLE_API_KEY"]:
            val = os.getenv(key)
            status = f"SET ({val[:6]}...)" if val else "NOT SET"
            print(f"  - {key:<20} : {status}")
        
        print("\n" + "=" * 50)
        all_ok = all(p['status'] == 'OK' for p in providers) if providers else False
        if all_ok:
            print("✓ All checks passed. Ready to go!")
        else:
            print("⚠ Some issues found. See above for details.")

    elif args.mode == "test-traffic":
        import asyncio
        
        async def run_test_async():
            print(f"Starting long-running LLM traffic simulation ({args.duration}s)...")
            print("Tip: Run 'python main.py tui' in another terminal to see results.")
            
            # 1. Start Proxy (as a subprocess so we can pipe to it)
            # Proxy command: python main.py proxy python tests/dummy_server.py
            # 프로시 실행 명령어 구성
            proxy_cmd = [
                sys.executable, sys.argv[0], "proxy",
                "--provider", args.provider,
                "--model", args.model,
                sys.executable, "tests/dummy_server.py"
            ]
            
            proxy_proc = await asyncio.create_subprocess_exec(
                *proxy_cmd,
                stdin=asyncio.subprocess.PIPE,
                stdout=asyncio.subprocess.PIPE,
                stderr=sys.stderr
            )
            
            # 2. Start Traffic Generator (as a subprocess)
            gen_cmd = [sys.executable, "tests/traffic_generator.py", str(args.duration)]
            gen_proc = await asyncio.create_subprocess_exec(
                *gen_cmd,
                stdin=asyncio.subprocess.PIPE,
                stdout=asyncio.subprocess.PIPE,
                stderr=sys.stderr
            )
            
            # 3. Connect them bidirectionally
            # gen_proc.stdout -> proxy_proc.stdin
            # proxy_proc.stdout -> gen_proc.stdin
            
            async def forward(source_stdout, target_stdin, label):
                while True:
                    line = await source_stdout.readline()
                    if not line:
                        break
                    target_stdin.write(line)
                    await target_stdin.drain()
            
            try:
                # Run correlation tasks
                tasks = [
                    asyncio.create_task(forward(gen_proc.stdout, proxy_proc.stdin, "G2P")),
                    asyncio.create_task(forward(proxy_proc.stdout, gen_proc.stdin, "P2G"))
                ]
                
                # Wait for the traffic generator to finish (it has the duration)
                await gen_proc.wait()
                
                # Cleanup
                for t in tasks:
                    t.cancel()
                
                if proxy_proc.returncode is None:
                    proxy_proc.terminate()
                    await proxy_proc.wait()
                    
            except Exception as e:
                print(f"Error during test: {e}")
            finally:
                print("\nTest simulation finished.")

        try:
            asyncio.run(run_test_async())
        except KeyboardInterrupt:
            print("\nTest terminated by user.")

    print("\nCheck complete.\n")

if __name__ == "__main__":
    main()
