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

VERSION = "0.1.9"

def show_splash():
    splash = r"""
        ________________________________________________
       /  ____________________________________________  \
      |  /                                            \  |
      | |    _      _      __  __    _  _             | |
      | |   | |    | |    |  \/  |  | || |   ___      | |
      | |   | |    | |    | |\/| |  | || |_ / __|     | |
      | |   | |___ | |___ | |  | |  |__   _|\__ \     | |
      | |   |_____||_____||_|  |_|     |_|  |___/     | |
      | |                                              | |
      | |         [ MCP LENS DISCOVERY ]               | |
      |  \____________________________________________/  |
       \________________________________________________/
              \________________________________/
    """
    print(splash)
    time.sleep(0.3)


def show_help():
    help_text = f"""
MCP-Lens v{VERSION}  —  K9s for MCP

USAGE
  llm4s                              Launch TUI dashboard (default)
  llm4s setup                        Interactive setup wizard
  llm4s proxy -- <command> [args]    Intercept MCP server traffic
  llm4s proxy --name <server>        Run a registered server via proxy
  llm4s config                       Show registered API keys
  llm4s config --add <provider>      Add or update an API key
  llm4s add-server <name> ...        Register an MCP server
  llm4s list-servers                 Show registered servers
  llm4s inject                        Inject llm4s proxy into Claude Desktop config
  llm4s inject --dry-run              Preview changes without modifying config
  llm4s restore                       Restore original Claude Desktop config
  llm4s api                           Start FastAPI backend (port 8000)
  llm4s doctor                        System health check
  llm4s help                          Show this help

QUICK START
  Step 1 — Configure your API key:
    llm4s setup
    (or: llm4s config --add gemini)

  Step 2 — Start the proxy (in a separate terminal):
    llm4s proxy -- node /path/to/your/mcp-server/index.js
    llm4s proxy --name my-server      ← if you registered one

  Step 3 — Open the dashboard:
    llm4s

SUPPORTED PROVIDERS
  gemini / google        Google Gemini (GOOGLE_API_KEY)
  claude / anthropic     Anthropic Claude (ANTHROPIC_API_KEY)
  openai                 OpenAI (OPENAI_API_KEY)

TUI KEYBINDINGS
  Enter     Drill down (Provider → Model → Session → Logs)
  Esc       Go back
  /         Search
  r         Replay message
  e         Edit payload
  x         Export session to Markdown
  F5        Refresh
  ?         Show all keybindings
  q         Quit

EXAMPLES
  llm4s config --add gemini AIza...
  llm4s add-server my-bot --command node -- /projects/bot/index.js
  llm4s proxy -- python /path/to/mcp_server.py
  llm4s proxy --name my-bot
"""
    print(help_text)


def run_setup_wizard():
    """대화형 셋업 위자드"""
    print("\n" + "=" * 50)
    print("  MCP-Lens Setup Wizard")
    print("=" * 50)

    provider_map = {
        "1": ("Google",    "GOOGLE_API_KEY",    "gemini"),
        "2": ("Anthropic", "ANTHROPIC_API_KEY", "claude"),
        "3": ("OpenAI",    "OPENAI_API_KEY",    "openai"),
    }

    # ── Step 1: Provider 선택 ──────────────────────────────
    print("\nStep 1/3  Which AI provider are you using?\n")
    print("  1. Google Gemini")
    print("  2. Anthropic Claude")
    print("  3. OpenAI")
    print("  4. Skip  (I'll set environment variables manually)")

    choice = input("\nEnter choice [1-4]: ").strip()

    if choice in provider_map:
        real_name, env_key, alias = provider_map[choice]
        print(f"\nProvider: {real_name}")
        print(f"Env var : {env_key}")

        api_key = getpass.getpass(f"Enter your {real_name} API key (hidden): ").strip()
        if api_key:
            from app.core.config import register_provider
            register_provider(real_name, env_key, api_key=api_key)
            print(f"\n✓ API key saved for {real_name}.")
        else:
            print("\n⚠ No key entered — skipping provider registration.")
    elif choice == "4":
        print("\nSkipped. Set env vars before running the proxy:")
        print("  export GOOGLE_API_KEY=...")
        print("  export ANTHROPIC_API_KEY=...")
        print("  export OPENAI_API_KEY=...")
    else:
        print("\n⚠ Invalid choice — skipping provider setup.")

    # ── Step 2: MCP 서버 등록 ─────────────────────────────
    print("\n" + "-" * 50)
    print("Step 2/3  Register an MCP server? [y/N]: ", end="")
    if input().strip().lower() == "y":
        name = input("  Server name (e.g. my-bot): ").strip()
        command = input("  Command    (e.g. node):    ").strip()
        args_raw = input("  Arguments  (e.g. /path/to/server.js): ").strip()
        args = args_raw.split() if args_raw else []
        if name and command:
            register_server(name, command, args)
            print(f"\n✓ Server '{name}' registered.")
            print(f"  Run with: llm4s proxy --name {name}")
        else:
            print("\n⚠ Name or command missing — skipping server registration.")
    else:
        print("  Skipped.")

    # ── Step 3: 완료 안내 ─────────────────────────────────
    print("\n" + "=" * 50)
    print("Setup complete! Next steps:\n")
    print("  Terminal A (proxy)    → llm4s proxy -- node /path/to/server.js")
    print("  Terminal B (dashboard)→ llm4s")
    print("\nTip: Run 'llm4s doctor' to verify your configuration.")
    print("=" * 50 + "\n")


def run_proxy(target_command: list, name: str = None,
              provider: str = "Unknown", model: str = "DirectCommand"):
    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)
            servers = list_servers()
            if servers:
                print("\nRegistered servers:", file=sys.stderr)
                for s in servers:
                    print(f"  - {s['name']}", file=sys.stderr)
            else:
                print("\nNo servers registered yet. Run: llm4s setup", 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 setup", file=sys.stderr)
        sys.exit(1)

    provider_name = provider
    server_name = model if model != "DirectCommand" else (name or "DirectCommand")

    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):
    uvicorn.run("app.api.main:app", host=host, port=port, reload=True)


def handle_config(args):
    from app.core.config import register_provider, list_providers

    if args.add:
        provider_alias = args.add[0].lower()
        # KEY가 인자로 없으면 대화형 프롬프트
        if len(args.add) >= 2:
            api_key = args.add[1]
        else:
            api_key = getpass.getpass(f"Enter API key for '{args.add[0]}' (hidden): ").strip()
            if not api_key:
                print("No key entered. Aborted.")
                return

        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"✓ API key registered for {real_name}.")
        else:
            register_provider(args.add[0], f"{args.add[0].upper()}_API_KEY", api_key=api_key)
            print(f"✓ API key registered for custom provider '{args.add[0]}'.")

    # --list 또는 단독 실행 시 목록 출력
    if args.list or not args.add:
        providers = list_providers()
        if not providers:
            print("No API keys registered.\n")
            print("Quick setup:")
            print("  llm4s setup                     ← interactive wizard")
            print("  llm4s config --add gemini        ← prompted input")
            print("  llm4s config --add gemini MY_KEY ← direct input")
            print("\nOr set environment variables:")
            print("  export GOOGLE_API_KEY=...")
            print("  export OPENAI_API_KEY=...")
            print("  export ANTHROPIC_API_KEY=...")
        else:
            print(f"\n{'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']}")


def handle_doctor():
    from app.core.config import CONFIG_PATH, DB_PATH, list_providers, list_servers
    import os

    print(f"\n[{APP_NAME} DOCTOR — System Health Check]")
    print("=" * 50)
    print(f"  Python  : {sys.version.split()[0]}")
    print(f"  Platform: {sys.platform}")
    print(f"  Version : v{VERSION}")

    print(f"\n  Config  : {CONFIG_PATH}")
    print(f"            {'✓ 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}")
    print(f"            {'✓ exists' + db_size if db_exists else '✗ not found'}")

    providers = list_providers()
    print(f"\n  Providers ({len(providers)}):")
    if not providers:
        print("    (none — run: llm4s setup)")
    for p in providers:
        icon = "✓" if p["status"] == "OK" else "✗"
        print(f"    {icon} {p['name']:<15} {p['status']} (via {p.get('source','?')})")

    servers = list_servers()
    print(f"\n  Servers ({len(servers)}):")
    if not servers:
        print("    (none — run: llm4s setup)")
    for s in servers:
        print(f"    - {s['name']:<15} {s['command']} {' '.join(s['args'])}")

    print("\n  Environment:")
    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:<22}: {status}")

    print("\n" + "=" * 50)
    all_ok = providers and all(p["status"] == "OK" for p in providers)
    if all_ok:
        print("✓ All checks passed. Ready to go!\n")
    else:
        print("⚠ Some issues found. See above for details.")
        if not providers:
            print("  → Run 'llm4s setup' to get started.\n")


def handle_inject(args):
    """llm4s inject [--dry-run] 핸들러"""
    from app.core.inject import inject_all_servers, get_claude_config_path

    if args.dry_run:
        print("[DRY RUN] 실제 파일은 변경되지 않습니다.\n")

    try:
        result = inject_all_servers(dry_run=args.dry_run)
    except (FileNotFoundError, RuntimeError) as e:
        print(f"오류: {e}", file=sys.stderr)
        sys.exit(1)

    print(f"설정 파일: {result['config_path']}\n")

    if result["wrapped"]:
        print(f"주입 완료 ({len(result['wrapped'])}개):")
        for name in result["wrapped"]:
            print(f"  + {name}")
    else:
        print("주입할 서버가 없습니다.")

    if result["skipped"]:
        print(f"\n건너뜀 (이미 주입됨, {len(result['skipped'])}개):")
        for name in result["skipped"]:
            print(f"  - {name}")

    if not args.dry_run and result["wrapped"]:
        print(f"\n백업 위치: {result['backup_path']}")
        print("\nClaude Desktop을 재시작하면 트래픽이 llm4s를 통해 흐릅니다.")
        print("모니터링 시작: llm4s")
        print("복원하려면:   llm4s restore")

    if args.dry_run and result.get("preview"):
        print("\n[변경 후 예상 mcpServers]:")
        import json as _json
        print(_json.dumps(result["preview"], indent=2, ensure_ascii=False))


def handle_restore():
    """llm4s restore 핸들러"""
    from app.core.inject import restore_all_servers

    try:
        result = restore_all_servers()
        print(f"복원 완료: {result['backup_path']}")
        print(f"         → {result['config_path']}")
        print("\nClaude Desktop을 재시작하세요.")
    except FileNotFoundError as e:
        print(f"오류: {e}", file=sys.stderr)
        sys.exit(1)


def main():
    show_splash()

    parser = argparse.ArgumentParser(
        prog="llm4s",
        description=f"MCP-Lens v{VERSION} — K9s for MCP",
        add_help=True,
    )
    parser.add_argument("--version", "-v", action="version", version=f"llm4s {VERSION}")
    subparsers = parser.add_subparsers(dest="mode", required=False)

    # 1. proxy
    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="MCP server command to intercept")

    # 2. tui
    subparsers.add_parser("tui", help="Open the TUI dashboard (default)")

    # 3. setup wizard
    subparsers.add_parser("setup", help="Interactive setup wizard (API keys + server registration)")

    # 4. config
    config_parser = subparsers.add_parser("config", help="Manage API keys / providers")
    config_parser.add_argument(
        "--add", nargs="+", metavar=("PROVIDER", "KEY"),
        help="Add API key: --add gemini [KEY]  (prompts if KEY omitted)"
    )
    config_parser.add_argument("--list", action="store_true", help="List registered providers")

    # 5. 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")

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

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

    # 8. doctor
    subparsers.add_parser("doctor", help="System health check")

    # 9. help (explicit subcommand)
    subparsers.add_parser("help", help="Show detailed help and keybindings")

    # 10. inject — Claude Desktop config 자동 주입
    inject_parser = subparsers.add_parser(
        "inject",
        help="Inject llm4s proxy into Claude Desktop MCP server config",
    )
    inject_parser.add_argument(
        "--dry-run", action="store_true",
        help="Preview changes without modifying the config file",
    )

    # 11. restore — inject 이전 상태 복원
    subparsers.add_parser(
        "restore",
        help="Restore Claude Desktop config to the state before llm4s inject",
    )

    # 12. test-traffic (internal dev tool)
    test_parser = subparsers.add_parser("test-traffic", help="Run LLM traffic simulation (dev)")
    test_parser.add_argument("--duration", type=int, default=10)
    test_parser.add_argument("--provider", default="Unknown")
    test_parser.add_argument("--model", default="DirectCommand")

    args = parser.parse_args()

    # ── Default: no subcommand → launch TUI ───────────────
    if args.mode is None or args.mode == "tui":
        from app.ui.app import run_tui
        run_tui()
        return

    if args.mode == "help":
        show_help()
        return

    if args.mode == "setup":
        run_setup_wizard()
        return

    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)
        return

    if args.mode == "add-server":
        args_list = args.args or []
        if args_list and args_list[0] == "--":
            args_list = args_list[1:]
        register_server(args.name, args.command, args_list)
        print(f"✓ Server '{args.name}' registered.")
        print(f"  Run with: llm4s proxy --name {args.name}")
        return

    if args.mode == "list-servers":
        servers = list_servers()
        if not servers:
            print("No servers registered.")
            print("  → llm4s setup     ← register one interactively")
            print("  → llm4s add-server <name> --command <cmd> -- <args>")
        else:
            print(f"\n{'NAME':<20} {'COMMAND':<15} ARGS")
            print("-" * 60)
            for s in servers:
                print(f"  {s['name']:<18} {s['command']:<15} {' '.join(s['args'])}")
        return

    if args.mode == "api":
        run_api(args.host, args.port)
        return

    if args.mode == "config":
        handle_config(args)
        return

    if args.mode == "doctor":
        handle_doctor()
        return

    if args.mode == "inject":
        handle_inject(args)
        return

    if args.mode == "restore":
        handle_restore()
        return

    if args.mode == "test-traffic":
        async def run_test_async():
            print(f"Starting LLM traffic simulation ({args.duration}s)...")
            print("Tip: open 'llm4s' in another terminal to watch.")

            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
            )
            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
            )

            async def forward(src, dst):
                while True:
                    line = await src.readline()
                    if not line:
                        break
                    dst.write(line)
                    await dst.drain()

            try:
                tasks = [
                    asyncio.create_task(forward(gen_proc.stdout, proxy_proc.stdin)),
                    asyncio.create_task(forward(proxy_proc.stdout, gen_proc.stdin)),
                ]
                await gen_proc.wait()
                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.")
        return


if __name__ == "__main__":
    main()
