#!/usr/bin/env python3
"""
gracestack CLI — DeepSeek V4 Pro on A100 GPUs, in your terminal.

Usage:
  gracestack chat "Explain quantum computing in one sentence"
  gracestack explain ./src/app.py
  gracestack key              # show/set API key
  gracestack status           # check your usage

Get a free API key: https://tools.gracestack.se/api.html
"""

import sys, os, json, urllib.request, urllib.error

API_URL = "https://tools.gracestack.se/api/chat"
CODE_URL = "https://tools.gracestack.se/api/code-explain"
KEY_URL = "https://tools.gracestack.se/api/register"
UPGRADE_URL = "https://tools.gracestack.se/power-pack.html"
CONFIG_DIR = os.path.expanduser("~/.config/gracestack")
KEY_FILE = os.path.join(CONFIG_DIR, "key")

def get_key():
    """Get API key from env or config file."""
    key = os.environ.get("GRACESTACK_KEY", "")
    if key:
        return key
    if os.path.isfile(KEY_FILE):
        return open(KEY_FILE).read().strip()
    return None

def save_key(key):
    os.makedirs(CONFIG_DIR, exist_ok=True)
    with open(KEY_FILE, "w") as f:
        f.write(key)
    os.chmod(KEY_FILE, 0o600)

def api_call(endpoint, payload, key):
    """Make an API call and return response."""
    data = json.dumps(payload).encode()
    url = endpoint
    req = urllib.request.Request(url, data=data, method="POST")
    req.add_header("Content-Type", "application/json")
    req.add_header("X-API-Key", key)
    try:
        resp = urllib.request.urlopen(req, timeout=60)
        body = json.loads(resp.read())
        return body, resp.status
    except urllib.error.HTTPError as e:
        body = json.loads(e.read())
        return body, e.code

def show_upgrade_nudge():
    print()
    print("─" * 50)
    print("⚡ Rate limit reached.")
    print(f"   Get unlimited access: {UPGRADE_URL}")
    print("   $1 USD — 100 calls, 24h, no auto-renew.")
    print("─" * 50)

def cmd_chat(prompt, key):
    """One-shot chat with the AI."""
    payload = {"message": prompt}
    resp, status = api_call(API_URL, payload, key)
    
    if status == 200:
        reply = resp.get("response") or resp.get("reply", "")
        rl = resp.get("rate_limit", {})
        remaining = rl.get("calls_remaining", "?")
        print(reply)
        print(f"\n[{remaining} calls remaining today]")
    elif status == 402:
        print("❌ Daily limit reached (1 free call/day).")
        show_upgrade_nudge()
    elif status == 401:
        print("❌ Invalid API key. Get one at: https://tools.gracestack.se/api.html")
    else:
        print(f"❌ Error {status}: {resp.get('detail', resp)}")

def cmd_explain(filepath, key):
    """Explain a code file."""
    if not os.path.isfile(filepath):
        print(f"❌ File not found: {filepath}")
        sys.exit(1)
    
    code = open(filepath).read()
    if len(code) > 10000:
        print("⚠ File too large (>10KB). Showing first 10KB.")
        code = code[:10000]
    
    lang = os.path.splitext(filepath)[1].lstrip(".")
    payload = {"code": code, "language": lang or "auto"}
    resp, status = api_call(CODE_URL, payload, key)
    
    if status == 200:
        explanation = resp.get("explanation") or resp.get("response", "")
        print(f"📄 {filepath}")
        print("─" * 50)
        print(explanation)
        rl = resp.get("rate_limit", {})
        remaining = rl.get("calls_remaining", "?")
        print(f"\n[{remaining} code explanations remaining today]")
    elif status == 402:
        print("❌ Daily limit reached.")
        show_upgrade_nudge()
    else:
        print(f"❌ Error {status}: {resp.get('detail', resp)}")

def cmd_key():
    """Show or set API key."""
    key = get_key()
    if key:
        print(f"🔑 API key: {key[:12]}...{key[-4:]}")
        print(f"   Stored in: {KEY_FILE}")
        print(f"   Or env: GRACESTACK_KEY")
    else:
        print("🔑 No API key set.")
        print(f"   Get one: {KEY_URL}")
        print(f"   Set it:  export GRACESTACK_KEY=your-key-here")
        print(f"   Or save: gracestack key YOUR-KEY")

def cmd_status(key):
    """Check API key status."""
    payload = {"message": "ping"}
    resp, status = api_call(API_URL, payload, key)
    
    if status == 200:
        rl = resp.get("rate_limit", {})
        remaining = rl.get("calls_remaining", "?")
        tier = rl.get("tier", "free")
        print(f"✅ API key valid")
        print(f"   Tier: {tier}")
        print(f"   Calls remaining today: {remaining}")
    elif status == 402:
        print(f"✅ API key valid (rate limited)")
        print(f"   Calls remaining: 0")
        print(f"   Upgrade: {UPGRADE_URL}")
    else:
        print("❌ API key invalid or expired.")

def main():
    if len(sys.argv) < 2:
        print(__doc__)
        sys.exit(0)
    
    cmd = sys.argv[1].lower()
    
    if cmd == "key":
        if len(sys.argv) > 2:
            save_key(sys.argv[2])
            print(f"🔑 API key saved to {KEY_FILE}")
        else:
            cmd_key()
        return
    
    key = get_key()
    if not key:
        print("🔑 No API key. Get one at: https://tools.gracestack.se/api.html")
        print(f"   Then run: gracestack key YOUR-KEY")
        print(f"   Or:      export GRACESTACK_KEY=your-key")
        sys.exit(1)
    
    if cmd == "chat":
        if len(sys.argv) < 3:
            print("Usage: gracestack chat \"your prompt here\"")
            sys.exit(1)
        cmd_chat(sys.argv[2], key)
    
    elif cmd in ("explain", "exp"):
        if len(sys.argv) < 3:
            print("Usage: gracestack explain ./src/app.py")
            sys.exit(1)
        cmd_explain(sys.argv[2], key)
    
    elif cmd == "status":
        cmd_status(key)
    
    elif cmd in ("code-lab", "codelab"):
        print(f"🌐 Code Lab: https://tools.gracestack.se/code-lab.html")
    
    elif cmd in ("upgrade", "powerpack", "power-pack"):
        print(f"⚡ Power Pack: {UPGRADE_URL}")
    
    else:
        print(f"Unknown command: {cmd}")
        print(__doc__)

if __name__ == "__main__":
    main()