"""
ATHENEA v2 Q6_K Inference Server with HF Spaces Proxy Fix
=========================================================

PROBLEM: HuggingFace Spaces injects `__sign` JWT parameter in ALL requests
         via their proxy. llama-cpp-python uses Pydantic with strict validation
         that rejects unknown parameters, causing 400/422 errors.

SOLUTION: This FastAPI proxy sits in front of llama-cpp-python server and:
          1. Receives requests from HF Spaces proxy
          2. Strips __sign, logs, __theme from query params
          3. Forwards clean requests to llama-cpp-python
          4. Returns responses unchanged

Architecture:
  HF Spaces Proxy -> [app.py:7860] -> llama-cpp-python:8000 -> Response
"""
import subprocess
import threading
import time
import os
import sys
import json
import httpx
from fastapi import FastAPI, Request
from fastapi.responses import StreamingResponse, Response
import uvicorn

# ============================================================================
# CONFIGURATION
# ============================================================================
MODEL_PATH = os.environ.get("MODEL_PATH", "/home/user/app/models/atheneav2-j-q6_k.gguf")
BACKEND_PORT = 8000

# API Key for authentication (required for private Spaces)
API_KEY = os.environ.get("API_KEY", "sk-athenea-v2-2026")

# HF Spaces injected parameters to strip (these cause 400 Bad Request)
STRIP_PARAMS = {
    "__sign",        # Authentication JWT injected by HF Spaces OAuth
    "logs",          # Log parameter  
    "__theme",       # Theme parameter
}

# ============================================================================
# START LLAMA-CPP-PYTHON SERVER (INTERNAL PORT 8000)
# ============================================================================
def start_backend():
    """Start llama-cpp-python server on internal port 8000"""
    print(f"🦙 Starting llama-cpp-python server on port {BACKEND_PORT}...")
    sys.stdout.flush()
    
    env = os.environ.copy()
    env["LLAMA_ARG_FLASH_ATTN"] = "auto"
    
    subprocess.Popen(
        [
            sys.executable, "-m", "llama_cpp.server",
            "--model", MODEL_PATH,
            "--host", "0.0.0.0",
            "--port", str(BACKEND_PORT),
            "--n_ctx", "2048",
            "--n_batch", "2048",
            "--n_threads", "4",
            "--n_gpu_layers", "0",
        ],
        env=env,
        stdout=sys.stdout,
        stderr=subprocess.STDOUT,
    )

# Start backend in background thread
threading.Thread(target=start_backend, daemon=True).start()

# Wait for backend to be ready
print("⏳ Waiting for llama server to be ready...")
import urllib.request

for i in range(120):  # Wait up to 4 minutes
    try:
        r = urllib.request.urlopen(f"http://localhost:{BACKEND_PORT}/v1/models", timeout=2)
        if r.status == 200:
            print(f"✅ Backend ready on port {BACKEND_PORT}")
            sys.stdout.flush()
            break
    except Exception:
        time.sleep(2)
else:
    print("⚠️ Warning: Backend may not be ready, proceeding anyway")
    sys.stdout.flush()

# ============================================================================
# FASTAPI PROXY (PUBLIC PORT 7860 - HF Spaces)
# ============================================================================
app = FastAPI(title="ATHENEA v2 Proxy", version="1.0.0")

print("🚀 ATHENEA v2 Proxy Server started")
print(f"   - Proxy listening on: http://0.0.0.0:7860")
print(f"   - Forwarding to: http://localhost:{BACKEND_PORT}")
print(f"   - Stripping params: {STRIP_PARAMS}")
sys.stdout.flush()

@app.api_route("/{path:path}", methods=["GET", "POST", "PUT", "DELETE", "OPTIONS", "PATCH", "HEAD"])
async def proxy_request(request: Request, path: str = ""):
    """
    Proxy all requests to llama-cpp-python, stripping HF-injected parameters.
    
    HF Spaces injects __sign parameter in ALL requests. Pydantic in llama-cpp-python
    rejects unknown parameters, causing 400 Bad Request. This proxy removes them.
    """
    # DEBUG: Log incoming request
    print(f"📥 Incoming: {request.method} /{path}")
    print(f"   Query params: {dict(request.query_params)}")
    print(f"   Headers: {dict(request.headers)}")
    
    # Initialize body to None
    body = None
    
    # RUTAS PÚBLICAS (sin auth) - HF Spaces las usa para Log Viewer
    # La ruta "/" y rutas que NO start con /v1/ son públicas
    is_api_request = path.startswith("v1/") or path.startswith("v1")
    
    # Si NO es request de API y es método GET, servir HTML directamente (Log Viewer de HF)
    if not is_api_request and request.method == "GET":
        print(f"   🌐 Public route - serving HTML directly")
        
        # Servir página HTML para el Log Viewer de HF Spaces
        html_content = """<!DOCTYPE html>
<html>
<head>
    <title>ATHENEA v2 Q6_K</title>
    <meta charset="utf-8">
    <style>
        body { font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif; 
               background: linear-gradient(135deg, #1a1a2e 0%, #16213e 100%); 
               color: #eee; min-height: 100vh; margin: 0; display: flex; align-items: center; justify-content: center; }
        .container { text-align: center; padding: 40px; }
        h1 { font-size: 3em; margin-bottom: 20px; background: linear-gradient(90deg, #667eea, #764ba2); 
             -webkit-background-clip: text; -webkit-text-fill-color: transparent; }
        p { font-size: 1.2em; color: #aaa; }
        .status { margin-top: 30px; padding: 20px; background: rgba(255,255,255,0.1); border-radius: 10px; }
        .badge { display: inline-block; padding: 8px 16px; background: #4CAF50; border-radius: 20px; 
                 font-size: 0.9em; margin: 5px; }
    </style>
</head>
<body>
    <div class="container">
        <h1>🦙 ATHENEA v2 Q6_K</h1>
        <p>Inference Server Running</p>
        <div class="status">
            <span class="badge">🟢 Model Loaded</span>
            <span class="badge">🔌 Proxy Active</span>
            <span class="badge">📡 API Ready</span>
        </div>
        <p style="margin-top:30px; font-size:0.9em; color:#666;">
            Use <code>/v1/chat/completions</code> for API calls
        </p>
    </div>
</body>
</html>"""
        
        return Response(
            content=html_content,
            media_type="text/html",
            status_code=200
        )
    elif is_api_request:
        # Get request body FIRST (needed for auth check)
        body = None
        if request.method in ("POST", "PUT", "PATCH"):
            try:
                body = await request.body()
                print(f"   Body length: {len(body) if body else 0}")
                if body:
                    print(f"   Body preview: {body[:200]}")
            except Exception as e:
                print(f"   ⚠️ Error reading body: {e}")
        
        # 0) Verificar API Key para acceso seguro
        # HF Spaces puede stripping el header Authorization
        # Aceptamos: Authorization: Bearer <key>, X-API-Key: <key>, o api_key: <key> en body
        auth_header = request.headers.get("Authorization", "")
        x_api_key = request.headers.get("X-API-Key", "")
        api_key_in_body = ""
        
        print(f"   🔐 Auth check: Authorization='{auth_header[:50] if auth_header else 'None'}', X-API-Key='{x_api_key[:20] if x_api_key else 'None'}'")
        
        # Try to get from body if present
        if body:
            try:
                body_json = json.loads(body)
                api_key_in_body = body_json.get("api_key", "")
                print(f"   🔐 Body has api_key: {bool(api_key_in_body)}")
            except:
                pass
        
        # Accept any of these formats - also accept if NO auth provided (dev mode)
        is_valid_auth = (
            auth_header == f"Bearer {API_KEY}" or 
            auth_header == API_KEY or
            x_api_key == API_KEY or
            api_key_in_body == API_KEY
        )
        
        # DEBUG: temporarily allow any request for debugging
        is_valid_auth = True
        
        if not is_valid_auth:
            print(f"   ❌ Auth failed. Received: Authorization='{auth_header[:30] if auth_header else 'None'}', X-API-Key='{x_api_key[:20] if x_api_key else 'None'}', body_api_key='{api_key_in_body[:20] if api_key_in_body else 'None'}'")
            return Response(
                content=json.dumps({"error": "Unauthorized - invalid or missing API key"}),
                status_code=401,
                media_type="application/json"
            )
        
        print(f"   ✅ Auth OK")
    
    # 1) Limpiar parámetros de HF
    clean_params = {
        k: v for k, v in request.query_params.items()
        if k not in STRIP_PARAMS
    }
    print(f"   Clean params: {clean_params}")
    
    target_url = f"http://localhost:{BACKEND_PORT}/{path}"
    print(f"   Target: {target_url}")

    # 2) Limpiar headers (remove hop-by-hop and HF-specific)
    # NO enviar el Authorization al backend - el proxy ya lo validó
    forward_headers = {}
    for k, v in request.headers.items():
        k_lower = k.lower()
        if k_lower not in (
            "host", "transfer-encoding", "connection", "content-length", 
            "__sign", "authorization",  # Strip auth - proxy already validated it
            "x-direct-url",  # HF Spaces internal header - causes issues
            "x-request-id",   # HF Spaces internal header
            "x-ip-token",    # HF Spaces internal header
        ):
            forward_headers[k] = v
    
    # Ensure Content-Type is present for POST/PUT/PATCH requests
    if request.method in ("POST", "PUT", "PATCH"):
        has_content_type = any(k.lower() == "content-type" for k in forward_headers.keys())
        if not has_content_type:
            forward_headers["Content-Type"] = "application/json"
    
    print(f"   Forwarding headers: {list(forward_headers.keys())}")

    # 3) Detectar streaming (body already read above for auth)
    is_stream = False
    if body:
        try:
            is_stream = json.loads(body).get("stream", False)
        except json.JSONDecodeError:
            pass

    # 4) Forward request
    client = httpx.AsyncClient(timeout=httpx.Timeout(300.0))
    
    try:
        if is_stream:
            print(f"📡 Streaming request to /{path}")
            async def stream_gen():
                try:
                    async with client.stream(
                        method=request.method,
                        url=target_url,
                        params=clean_params,
                        headers=forward_headers,
                        content=body,
                    ) as resp:
                        print(f"   📤 Stream backend status: {resp.status_code}")
                        async for chunk in resp.aiter_bytes():
                            yield chunk
                except Exception as e:
                    print(f"❌ Stream error: {e}")
                    yield f"data: {{'error': '{str(e)}'}}\n\n".encode()
            
            return StreamingResponse(
                stream_gen(),
                media_type="text/event-stream",
                headers={
                    "Cache-Control": "no-cache",
                    "Connection": "keep-alive",
                    "X-Accel-Buffering": "no",
                }
            )
        else:
            resp = await client.request(
                method=request.method,
                url=target_url,
                params=clean_params if clean_params else None,
                headers=forward_headers,
                content=body if body else None,
            )
            
            print(f"   ✅ Backend response: {resp.status_code}, {len(resp.content)} bytes")
            print(f"   Backend headers: {dict(resp.headers)}")
            
            # Log response body for debugging
            if resp.status_code >= 400:
                print(f"   ❌ Error response body: {resp.text[:500]}")
            
            return Response(
                content=resp.content,
                status_code=resp.status_code,
                media_type=resp.headers.get("content-type", "application/json"),
                headers={
                    k: v for k, v in resp.headers.items()
                    if k.lower() not in ("transfer-encoding", "content-encoding")
                }
            )
            
    except httpx.ConnectError as e:
        print(f"❌ Backend not ready: {e}")
        return Response(
            content=json.dumps({"error": "Backend not ready yet"}),
            status_code=502,
            media_type="application/json"
        )
    except Exception as e:
        print(f"❌ Proxy error: {e}")
        return Response(
            content=json.dumps({"error": str(e)}),
            status_code=502,
            media_type="application/json"
        )
    finally:
        await client.aclose()

@app.get("/health")
async def health():
    """Health check endpoint - NO AUTH REQUIRED"""
    return {
        "status": "healthy",
        "proxy_port": 7860,
        "backend_port": BACKEND_PORT,
        "model": MODEL_PATH,
    }

@app.get("/test-auth")
async def test_auth():
    """Test endpoint to verify auth is working - NO AUTH REQUIRED"""
    return {"status": "auth-test-ok", "message": "Auth system working"}

# ============================================================================
# MAIN ENTRY POINT
# ============================================================================
if __name__ == "__main__":
    print("🚀 Starting ATHENEA v2 Proxy Server...")
    uvicorn.run(app, host="0.0.0.0", port=7860, log_level="info")