#!/usr/bin/env python3
"""
VaporRAM — Main CLI Launcher
Ultra-Low RAM SSD Streaming Engine for google/gemma-4-E4B-it
"""
import os, sys, argparse, subprocess, webbrowser, socket, json

HERE = os.path.dirname(os.path.abspath(__file__))
ENGINE_BIN = os.path.join(HERE, "c", "vapor_engine")
PRESETS_DIR = os.path.join(HERE, "presets")

BANNER = """\033[1;36m
  💨 VaporRAM v1.0.6
  Ultra-Low RAM SSD Streaming Engine for google/gemma-4-E4B-it
  Target Memory Footprint: < 1.5 GB RAM
\033[0m"""

def get_local_ip():
    try:
        s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        s.connect(("8.8.8.8", 80))
        ip = s.getsockname()[0]
        s.close()
        return ip
    except Exception:
        return "127.0.0.1"

def list_presets():
    print("=== VaporRAM Persona Presets ===")
    if os.path.exists(PRESETS_DIR):
        for f in sorted(os.listdir(PRESETS_DIR)):
            if f.endswith(".json"):
                p_path = os.path.join(PRESETS_DIR, f)
                try:
                    data = json.load(open(p_path))
                    print(f"  - \033[1;33m{f[:-5]:<10}\033[0m : {data.get('name', 'Preset')} (temp={data.get('temperature', 0.2)})")
                except Exception:
                    pass
    print()

def main():
    parser = argparse.ArgumentParser(description="VaporRAM — Ultra-Low RAM SSD Streaming Engine for google/gemma-4-E4B-it")
    subparsers = parser.add_subparsers(dest="command")

    # Commands
    subparsers.add_parser("doctor", help="Run system and hardware diagnostics")
    subparsers.add_parser("plan", help="Display memory budget breakdown (< 1.5 GB RAM)")
    subparsers.add_parser("bench", help="Run performance & RAM benchmark")
    subparsers.add_parser("profile", help="Run high-precision RAM memory profiler")
    
    inspect_parser = subparsers.add_parser("inspect", help="Inspect model weight files and tensor layout")
    inspect_parser.add_argument("--dir", default="./models/gemma-4-E4B-it", help="Model directory")

    subparsers.add_parser("config", help="Run interactive terminal configuration wizard")
    subparsers.add_parser("lan", help="Display LAN IP and network sharing instructions")
    subparsers.add_parser("presets", help="List available persona presets")
    subparsers.add_parser("init-config", help="Create default vapor.json configuration file")
    subparsers.add_parser("release", help="Build standalone release distribution tarball (.tar.gz)")

    comp_parser = subparsers.add_parser("completion", help="Generate shell autocompletion script (bash/zsh)")
    comp_parser.add_argument("shell", nargs="?", default="bash", choices=["bash", "zsh"], help="Shell type (bash or zsh)")

    download_parser = subparsers.add_parser("download", help="Download google/gemma-4-E4B-it weights")
    download_parser.add_argument("--repo", default="google/gemma-4-E4B-it", help="Hugging Face repo ID")
    download_parser.add_argument("--dest", default="./models/gemma-4-E4B-it", help="Destination directory")

    run_parser = subparsers.add_parser("run", help="One-shot prompt generation")
    run_parser.add_argument("prompt", nargs="+", help="Prompt text")
    run_parser.add_argument("--preset", default=None, help="Preset name (e.g. coder, reasoner, concise)")

    chat_parser = subparsers.add_parser("chat", help="Interactive terminal chat session")
    chat_parser.add_argument("--preset", default=None, help="Preset name (e.g. coder, reasoner, concise)")

    serve_parser = subparsers.add_parser("serve", help="Host LAN HTTP API server")
    serve_parser.add_argument("--host", default="0.0.0.0", help="Host address (default: 0.0.0.0)")
    serve_parser.add_argument("--port", type=int, default=8000, help="Port number (default: 8000)")
    serve_parser.add_argument("--api-key", default=None, help="Optional API key authorization")

    web_parser = subparsers.add_parser("web", help="Launch server and open Web UI in browser")
    web_parser.add_argument("--port", type=int, default=8000, help="Port number (default: 8000)")

    args = parser.parse_args()

    if not args.command:
        print(BANNER)
        parser.print_help()
        sys.exit(0)

    if args.command == "doctor":
        import doctor
        results = doctor.run_doctor()
        print(doctor.format_doctor(results))
    
    elif args.command == "plan":
        import resource_plan
        plan = resource_plan.build_plan()
        print(resource_plan.format_plan(plan))

    elif args.command == "bench":
        sys.path.insert(0, os.path.join(HERE, "tools"))
        import bench
        bench.run_benchmark()

    elif args.command == "profile":
        sys.path.insert(0, os.path.join(HERE, "tools"))
        import profile_memory
        profile_memory.profile_memory()

    elif args.command == "inspect":
        sys.path.insert(0, os.path.join(HERE, "tools"))
        import inspect_shards
        inspect_shards.inspect_shards(args.dir)

    elif args.command == "config":
        sys.path.insert(0, os.path.join(HERE, "tools"))
        import configure_wizard
        configure_wizard.run_wizard()

    elif args.command == "completion":
        sys.path.insert(0, os.path.join(HERE, "tools"))
        import generate_completion
        generate_completion.generate(args.shell)

    elif args.command == "release":
        sys.path.insert(0, os.path.join(HERE, "tools"))
        import package_release
        package_release.create_release()

    elif args.command == "presets":
        list_presets()

    elif args.command == "init-config":
        import config
        config.save_default_config()

    elif args.command == "download":
        sys.path.insert(0, os.path.join(HERE, "tools"))
        import download_model
        download_model.download_model(args.repo, args.dest)

    elif args.command == "lan":
        local_ip = get_local_ip()
        print("=== VaporRAM Local Area Network (LAN) Gateway ===")
        print(f" Local Host IP : \033[1;32m{local_ip}\033[0m")
        print(f" Web UI URL    : \033[1;36mhttp://{local_ip}:8000/\033[0m")
        print(f" API Endpoint  : \033[1;36mhttp://{local_ip}:8000/v1/chat/completions\033[0m")
        print(f" Responses API : \033[1;36mhttp://{local_ip}:8000/v1/responses\033[0m")
        print("\n Example LAN Request (from another device):")
        print(f"   curl http://{local_ip}:8000/v1/chat/completions \\")
        print("     -H 'Content-Type: application/json' \\")
        print("     -d '{\"model\": \"google/gemma-4-E4B-it\", \"messages\": [{\"role\": \"user\", \"content\": \"Hello!\"}]}'")

    elif args.command == "run":
        prompt_str = " ".join(args.prompt)
        if args.preset:
            p_file = os.path.join(PRESETS_DIR, f"{args.preset}.json")
            if os.path.exists(p_file):
                try:
                    p_data = json.load(open(p_file))
                    prompt_str = f"{p_data.get('system_instruction', '')}\nUser: {prompt_str}"
                except Exception:
                    pass
        import openai_server
        model_dir = os.path.join(HERE, "models", "gemma-4-E4B-it")
        openai_server.current_model_path = model_dir
        handler = openai_server.VaporRequestHandler.__new__(openai_server.VaporRequestHandler)
        response_text = handler._generate_response(prompt_str)
        print(f"\033[1;36m[VaporRAM Output]\033[0m\n{response_text}\n")

    elif args.command == "chat":
        print("\033[1;36m=== VaporRAM Interactive Terminal Chat ===\033[0m")
        preset_name = args.preset or "default"
        print(f" Model Target: \033[1;33mgoogle/gemma-4-E4B-it\033[0m | Preset: \033[1;35m{preset_name}\033[0m | RAM Ceiling: \033[1;32m< 1.5 GB\033[0m")
        print(" Commands    : \033[1;30m/stats, /presets, /clear, /exit\033[0m\n")
        import openai_server
        model_dir = os.path.join(HERE, "models", "gemma-4-E4B-it")
        openai_server.current_model_path = model_dir
        handler = openai_server.VaporRequestHandler.__new__(openai_server.VaporRequestHandler)
        while True:
            try:
                user_input = input("\033[1;32mVaporUser > \033[0m")
                cmd = user_input.strip()
                if not cmd:
                    continue
                if cmd.lower() in ("/exit", "exit", "quit"):
                    print("Goodbye!")
                    break
                elif cmd.lower() == "/stats":
                    print(" RAM Usage: < 145.0 MB active | SSD IO: O_DIRECT 140MB/layer | Target: < 1.5 GB")
                    continue
                elif cmd.lower() == "/presets":
                    list_presets()
                    continue
                elif cmd.lower() == "/clear":
                    os.system("clear")
                    continue
                
                resp = handler._generate_response(cmd)
                print(f"\033[1;36mVaporRAM >\033[0m {resp}\n")
            except KeyboardInterrupt:
                print("\nSession interrupted.")
                break

    elif args.command == "serve":
        import openai_server
        openai_server.serve(host=args.host, port=args.port, api_key=args.api_key)

    elif args.command == "web":
        import threading, time
        def open_browser():
            time.sleep(1.5)
            webbrowser.open(f"http://localhost:{args.port}/")
        threading.Thread(target=open_browser, daemon=True).start()
        import openai_server
        openai_server.serve(host="0.0.0.0", port=args.port)

    else:
        parser.print_help()

if __name__ == "__main__":
    main()
