"""
TokioAI Engine -- CLI Entry Point

Usage:
  python -m engine serve              # Start OpenAI-compatible API server
  python -m engine test               # Run test suite
  python -m engine bench              # Benchmark pipeline
  python -m engine chat               # Interactive chat (for testing)
  python -m engine stats              # Show memory/pipeline stats
"""

import sys
import json
import time
import argparse

from .pipeline import Pipeline, EngineConfig


def cmd_serve(args):
    """Start the API server."""
    from .server import create_app
    from .pipeline import EngineConfig

    config = EngineConfig(
        ollama_host=args.ollama_host,
        model_name=args.model,
        max_retries=args.max_retries,
        use_guided_gen=not args.no_guided,
        use_memory=not args.no_memory,
        verbose=args.verbose,
    )

    print(f"""
 _____     _   _       _    ___
|_   _|__ | | _(_) ___ / \\  |_ _|
  | |/ _ \\| |/ / |/ _ / _ \\  | |
  | | (_) |   <| | (_/ ___ \\ | |
  |_|\\___/|_|\\_\\_|\\___/_/   \\_\\___|
   ___           _
  | __| _ _  __ _(_) _ _  ___
  | _| | ' \\/ _` | || ' \\/ -_)
  |___||_||_\\__, |_||_||_\\___|
            |___/            v3.0.0

  Model:    {config.model_name}
  Ollama:   {config.ollama_host}
  Memory:   {'ON' if config.use_memory else 'OFF'}
  Guided:   {'ON' if config.use_guided_gen else 'OFF'}
  Retries:  {config.max_retries}
  
  API: http://{args.host}:{args.port}/v1/chat/completions
  Health: http://{args.host}:{args.port}/health
  Stats: http://{args.host}:{args.port}/stats
""")

    app = create_app(config)
    app.run(host=args.host, port=args.port, debug=False)


def cmd_chat(args):
    """Interactive chat for testing the pipeline."""
    config = EngineConfig(
        ollama_host=args.ollama_host,
        model_name=args.model,
        verbose=args.verbose,
    )
    engine = Pipeline(config)

    print(f"\n[TokioAI Engine] Connected to {config.ollama_host} / {config.model_name}")
    print("[TokioAI Engine] Type 'quit' to exit, 'stats' for stats, 'reset' to clear history\n")

    while True:
        try:
            user_input = input("\033[36myou>\033[0m ")
        except (EOFError, KeyboardInterrupt):
            print("\nbye.")
            break

        if not user_input.strip():
            continue
        if user_input.strip().lower() == "quit":
            break
        if user_input.strip().lower() == "stats":
            print(json.dumps(engine.get_stats(), indent=2))
            continue
        if user_input.strip().lower() == "reset":
            engine.reset()
            print("[reset]")
            continue

        result = engine.process(user_input)

        # Display result
        print(f"\033[33m[{result.intent}|{result.confidence:.1f}|{result.attempts}att|{result.latency_ms}ms]\033[0m")

        if result.output_type == "tool_call":
            for tc in result.tool_calls:
                print(f"\033[32m>>> TOOL: {tc['name']}({json.dumps(tc['arguments'])})\033[0m")
        else:
            print(f"\033[37m{result.text}\033[0m")
        print()


def cmd_test(args):
    """Run the test suite."""
    from .test_engine import run_tests
    run_tests(verbose=args.verbose, ollama_host=args.ollama_host, model=args.model)


def cmd_bench(args):
    """Benchmark the pipeline."""
    config = EngineConfig(
        ollama_host=args.ollama_host,
        model_name=args.model,
        verbose=False,
    )
    engine = Pipeline(config)

    test_prompts = [
        ("check disk space", "tool_call"),
        ("what is docker", "text"),
        ("read /etc/hostname", "tool_call"),
        ("hi", "text"),
        ("find all python files with TODO", "tool_call"),
        ("explain SQL injection", "text"),
        ("is docker running", "tool_call"),
        ("list processes by memory usage", "tool_call"),
        ("what is the TCP 3-way handshake", "text"),
        ("scan ports on localhost", "tool_call"),
    ]

    print(f"\n[BENCH] Running {len(test_prompts)} prompts against {config.model_name}...\n")
    
    results = []
    for prompt, expected in test_prompts:
        start = time.time()
        result = engine.process(prompt)
        elapsed = int((time.time() - start) * 1000)

        correct = (
            (expected == "tool_call" and result.output_type == "tool_call") or
            (expected == "text" and result.output_type == "text")
        )

        status = "\033[32mPASS\033[0m" if correct else "\033[31mFAIL\033[0m"
        detail = ""
        if result.output_type == "tool_call" and result.tool_calls:
            detail = f"-> {result.tool_calls[0]['name']}"
        elif result.output_type == "text":
            detail = f"-> text ({len(result.text)} chars)"

        print(f"  {status} [{elapsed:4d}ms] [{result.intent:9s}] {prompt:40s} {detail}")
        results.append(correct)
        engine.reset()  # Reset between tests

    passed = sum(results)
    total = len(results)
    print(f"\n[BENCH] {passed}/{total} correct ({100*passed//total}%)")
    print(f"[BENCH] Stats: {json.dumps(engine.get_stats(), indent=2)}")


def cmd_stats(args):
    """Show memory and pipeline stats."""
    from .memory import EngineMemory
    mem = EngineMemory()
    print(json.dumps(mem.get_stats(), indent=2))


def main():
    parser = argparse.ArgumentParser(description="TokioAI Engine")
    parser.add_argument("--ollama-host", default=None, help="Ollama API URL")
    parser.add_argument("--model", default="tokioai-coder", help="Model name")
    parser.add_argument("-v", "--verbose", action="store_true")
    
    subparsers = parser.add_subparsers(dest="command")
    
    # serve
    serve_parser = subparsers.add_parser("serve", help="Start API server")
    serve_parser.add_argument("--host", default="0.0.0.0", help="Bind address")
    serve_parser.add_argument("--port", type=int, default=8080, help="Port")
    serve_parser.add_argument("--ollama-host", default=None, help="Ollama API URL")
    serve_parser.add_argument("--model", default="tokioai-coder", help="Model name")
    serve_parser.add_argument("--no-memory", action="store_true", help="Disable memory")
    serve_parser.add_argument("--no-guided", action="store_true", help="Disable guided gen")
    serve_parser.add_argument("--max-retries", type=int, default=2, help="Max retries")
    serve_parser.add_argument("-v", "--verbose", action="store_true")
    
    # chat
    chat_parser = subparsers.add_parser("chat", help="Interactive chat")
    
    # test
    test_parser = subparsers.add_parser("test", help="Run tests")
    
    # bench
    bench_parser = subparsers.add_parser("bench", help="Benchmark")
    
    # stats
    stats_parser = subparsers.add_parser("stats", help="Show stats")

    # Add common args to all subparsers that need them
    for sp in [chat_parser, test_parser, bench_parser]:
        sp.add_argument("--ollama-host", default=None, help="Ollama API URL")
        sp.add_argument("--model", default="tokioai-coder", help="Model name")
        sp.add_argument("-v", "--verbose", action="store_true")

    args = parser.parse_args()

    if args.command == "serve":
        cmd_serve(args)
    elif args.command == "chat":
        cmd_chat(args)
    elif args.command == "test":
        cmd_test(args)
    elif args.command == "bench":
        cmd_bench(args)
    elif args.command == "stats":
        cmd_stats(args)
    else:
        parser.print_help()


if __name__ == "__main__":
    main()