"""sdkspindle trial proxy — Vercel Edge Function.

Deploy:
  1. cd proxy/
  2. vercel deploy
  3. Set GOOGLE_API_KEY in Vercel dashboard env vars

Rate limit: 20 queries/day per IP (stored in Vercel KV or in-memory fallback).

Endpoint: POST /api/query
  Body: {"question": "...", "context": "...", "model": "gemini-2.0-flash"}
  Returns: {"answer": "..."} or {"error": "..."}
"""

import json
import os
import time
import hashlib
from http.server import BaseHTTPRequestHandler

# In-memory rate limiter (resets on cold start, ~5min on Vercel)
# For production, swap to Vercel KV or Upstash Redis.
_rate_store = {}  # ip_hash -> {count, window_start}
RATE_LIMIT = 20  # queries per day
RATE_WINDOW = 86400  # 24 hours in seconds

SYSTEM_PROMPT = """You are sdkspindle, an expert on the com.google.genai Java SDK (google-genai).
You have access to the extracted SDK API reference below.

Rules:
- Answer with EXACT Java/Kotlin types, method signatures, and field names from the reference.
- When showing code, use Kotlin syntax (the primary consumer language).
- Always show the correct import statements.
- If a field returns Optional<T>, always show the .orElse() unwrapping.
- If a builder method has varargs overloads, mention both List<T> and varargs options.
- Be concise. Developers want the answer, not an essay.
- If the question is about something not in the reference, say so clearly."""

ALLOWED_ORIGINS = ["*"]  # Lock down for production
MAX_CONTEXT_SIZE = 50_000  # 50KB context cap
MAX_QUESTION_LEN = 500


def _hash_ip(ip: str) -> str:
    return hashlib.sha256(ip.encode()).hexdigest()[:16]


def _check_rate_limit(ip: str) -> tuple:
    """Returns (allowed: bool, remaining: int, reset_at: int)."""
    now = time.time()
    ip_hash = _hash_ip(ip)
    entry = _rate_store.get(ip_hash)

    if not entry or now - entry["window_start"] > RATE_WINDOW:
        _rate_store[ip_hash] = {"count": 1, "window_start": now}
        return True, RATE_LIMIT - 1, int(now + RATE_WINDOW)

    if entry["count"] >= RATE_LIMIT:
        return False, 0, int(entry["window_start"] + RATE_WINDOW)

    entry["count"] += 1
    remaining = RATE_LIMIT - entry["count"]
    return True, remaining, int(entry["window_start"] + RATE_WINDOW)


def _call_gemini(question: str, context: str, model: str) -> str:
    """Call Gemini API server-side."""
    api_key = os.environ.get("GOOGLE_API_KEY")
    if not api_key:
        raise RuntimeError("Server misconfigured: no GOOGLE_API_KEY")

    import urllib.request
    import urllib.error

    url = f"https://generativelanguage.googleapis.com/v1beta/models/{model}:generateContent?key={api_key}"

    payload = json.dumps({
        "system_instruction": {"parts": [{"text": SYSTEM_PROMPT}]},
        "contents": [{"parts": [{"text": f"SDK API Reference (relevant subset):\n```json\n{context}\n```\n\nQuestion: {question}"}]}],
        "generationConfig": {
            "temperature": 0.1,
            "maxOutputTokens": 2048,
        }
    }).encode("utf-8")

    req = urllib.request.Request(
        url,
        data=payload,
        headers={"Content-Type": "application/json"},
        method="POST",
    )

    with urllib.request.urlopen(req, timeout=30) as resp:
        data = json.loads(resp.read())

    # Extract text from Gemini response
    candidates = data.get("candidates", [])
    if not candidates:
        return "(no response from Gemini)"
    parts = candidates[0].get("content", {}).get("parts", [])
    return "".join(p.get("text", "") for p in parts) or "(empty response)"


class handler(BaseHTTPRequestHandler):
    """Vercel serverless function handler."""

    def do_POST(self):
        # CORS
        self.send_header("Access-Control-Allow-Origin", "*")
        self.send_header("Access-Control-Allow-Methods", "POST, OPTIONS")
        self.send_header("Access-Control-Allow-Headers", "Content-Type")

        # Rate limit
        client_ip = self.headers.get("x-forwarded-for", self.client_address[0]).split(",")[0].strip()
        allowed, remaining, reset_at = _check_rate_limit(client_ip)

        if not allowed:
            self.send_response(429)
            self.send_header("Content-Type", "application/json")
            self.send_header("X-RateLimit-Remaining", str(remaining))
            self.send_header("X-RateLimit-Reset", str(reset_at))
            self.end_headers()
            self.wfile.write(json.dumps({
                "error": f"Rate limit exceeded ({RATE_LIMIT}/day). Set your own GOOGLE_API_KEY for unlimited.",
                "remaining": remaining,
                "reset_at": reset_at,
            }).encode())
            return

        # Parse body
        content_length = int(self.headers.get("Content-Length", 0))
        if content_length > 100_000:  # 100KB body cap
            self.send_response(413)
            self.send_header("Content-Type", "application/json")
            self.end_headers()
            self.wfile.write(json.dumps({"error": "Request too large"}).encode())
            return

        body = self.rfile.read(content_length)
        try:
            data = json.loads(body)
        except json.JSONDecodeError:
            self.send_response(400)
            self.send_header("Content-Type", "application/json")
            self.end_headers()
            self.wfile.write(json.dumps({"error": "Invalid JSON"}).encode())
            return

        question = (data.get("question") or "")[:MAX_QUESTION_LEN]
        context = (data.get("context") or "")[:MAX_CONTEXT_SIZE]
        model = data.get("model", "gemini-2.0-flash")

        if not question:
            self.send_response(400)
            self.send_header("Content-Type", "application/json")
            self.end_headers()
            self.wfile.write(json.dumps({"error": "Missing 'question' field"}).encode())
            return

        # Validate model (whitelist)
        allowed_models = {
            "gemini-2.0-flash", "gemini-2.5-flash",
            "gemini-2.0-flash-lite", "gemini-2.5-flash-lite",
        }
        if model not in allowed_models:
            model = "gemini-2.0-flash"

        # Call Gemini
        try:
            answer = _call_gemini(question, context, model)
        except Exception as e:
            self.send_response(502)
            self.send_header("Content-Type", "application/json")
            self.end_headers()
            self.wfile.write(json.dumps({"error": f"Gemini API error: {str(e)[:200]}"}).encode())
            return

        self.send_response(200)
        self.send_header("Content-Type", "application/json")
        self.send_header("X-RateLimit-Remaining", str(remaining))
        self.send_header("X-RateLimit-Reset", str(reset_at))
        self.end_headers()
        self.wfile.write(json.dumps({
            "answer": answer,
            "remaining": remaining,
            "model": model,
        }).encode())

    def do_OPTIONS(self):
        self.send_response(200)
        self.send_header("Access-Control-Allow-Origin", "*")
        self.send_header("Access-Control-Allow-Methods", "POST, OPTIONS")
        self.send_header("Access-Control-Allow-Headers", "Content-Type")
        self.end_headers()
