#!/usr/bin/env python3
"""
Sahon AI - Gradio Bootstrap for Transformers.js (Node.js)
=========================================================
This app starts a Node.js child process that runs Transformers.js.
Gradio acts as the UI proxy layer.
"""

import os
import sys
import json
import time
import subprocess
import threading
import urllib.request
import urllib.error
import atexit
import signal

# ─── ZeroGPU Activation ───
# ZeroGPU requires at least one @spaces.GPU decorated function.
# Even if we don't use GPU (Transformers.js runs on CPU),
# the decorator must be present for ZeroGPU to activate.

try:
    from spaces import GPU as spaces_gpu

    @spaces_gpu
    def _zerogpu_placeholder():
        """Satisfy ZeroGPU requirement. GPU not actually used."""
        return True

    HAS_SPACES = True
    print("[Sahon] ✅ ZeroGPU compatible (placeholder registered)")
except ImportError:
    HAS_SPACES = False
    print("[Sahon] ⚠️ spaces module not available")

# ─── Config ───
NODE_SERVER_PORT = 8888
NODE_SERVER_URL = f"http://127.0.0.1:{NODE_SERVER_PORT}"
NODE_SCRIPT = "server.mjs"

# ─── Node.js Process Management ───
node_process = None

def start_node_server():
    """Start the Node.js Transformers.js server as a subprocess."""
    global node_process
    
    # Install npm packages first
    print("[Sahon] Installing npm packages...")
    npm_install = subprocess.run(
        ["npm", "install"],
        capture_output=True, text=True, timeout=120
    )
    if npm_install.returncode != 0:
        print(f"[Sahon] npm install stderr: {npm_install.stderr}")
    
    # Start Node.js server
    print(f"[Sahon] Starting Node.js server on port {NODE_SERVER_PORT}...")
    node_process = subprocess.Popen(
        ["node", NODE_SCRIPT],
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        text=True,
        bufsize=1,
    )
    
    # Log output in background
    def log_output():
        for line in node_process.stdout:
            print(f"[Node.js] {line}", end="")
    
    threading.Thread(target=log_output, daemon=True).start()
    
    print("[Sahon] Waiting for Node.js server to start...")
    time.sleep(2)
    
    # Wait for health check to pass
    for i in range(30):
        try:
            req = urllib.request.Request(f"{NODE_SERVER_URL}/health")
            resp = urllib.request.urlopen(req, timeout=5)
            data = json.loads(resp.read())
            if data.get("model_ready"):
                print("[Sahon] ✅ Node.js server ready!")
                return
            print(f"[Sahon]   Waiting for model... attempt {i+1}")
        except Exception as e:
            print(f"[Sahon]   Waiting for server... attempt {i+1} ({e})")
        time.sleep(5)
    
    print("[Sahon] ⚠️ Node.js server started but model may not be ready yet")

def stop_node_server():
    """Stop the Node.js process."""
    global node_process
    if node_process:
        print("[Sahon] Stopping Node.js server...")
        node_process.terminate()
        try:
            node_process.wait(timeout=10)
        except:
            node_process.kill()
        node_process = None

atexit.register(stop_node_server)

# ─── Gradio UI ───
import gradio as gr
from gradio.routes import App
from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse
import httpx

def chat_function(message: str, history: list) -> str:
    """Send chat request to Node.js Transformers.js server."""
    # Build messages from history
    messages = []
    for user_msg, assistant_msg in history:
        messages.append({"role": "user", "content": user_msg})
        if assistant_msg:
            messages.append({"role": "assistant", "content": assistant_msg})
    messages.append({"role": "user", "content": message})
    
    # Call Node.js server
    try:
        payload = json.dumps({
            "model": "phi-3-mini-4k-instruct",
            "messages": messages,
            "temperature": 0.7,
            "max_tokens": 512,
        }).encode()
        
        req = urllib.request.Request(
            f"{NODE_SERVER_URL}/v1/chat/completions",
            data=payload,
            headers={"Content-Type": "application/json"},
            method="POST"
        )
        resp = urllib.request.urlopen(req, timeout=120)
        data = json.loads(resp.read())
        
        content = data["choices"][0]["message"]["content"]
        
        # Add Mission Barisal quality info if available
        mission = data.get("_mission_barisal", {})
        if mission:
            quality = mission.get("quality_score", 0)
            content += f"\n\n---\n_Mission Barisal: {quality*100:.0f}% quality_"
        
        return content
    except urllib.error.HTTPError as e:
        return f"❌ Error {e.code}: {e.read().decode()[:200]}"
    except urllib.error.URLError as e:
        if "model_ready" not in str(e):
            return "⏳ Model is loading... Please wait and try again."
        return f"❌ Connection error: {e.reason}"
    except Exception as e:
        return f"❌ Error: {str(e)[:200]}"

def run_app():
    """Main app function — wrapped with @spaces.GPU for ZeroGPU."""
    
    # Start Node.js server in background
    threading.Thread(target=start_node_server, daemon=True).start()
    
    # Create Gradio UI
    with gr.Blocks(
        title="Sahon AI - Transformers.js + Mission Barisal",
        theme=gr.themes.Soft(),
    ) as demo:
        gr.Markdown("""
        # 🤖 Sahon AI
        ### Transformers.js (Node.js) + Gradio + ZeroGPU
        
        **JavaScript-powered LLM on Hugging Face Spaces!**  
        No Python ML dependencies — pure Transformers.js inference.
        """)
        
        with gr.Row():
            status_box = gr.Textbox(
                value="Starting Node.js server...",
                label="🟡 Model Status",
                interactive=False,
            )
        
        gr.ChatInterface(
            fn=chat_function,
            title="💬 Chat",
            description="Powered by Xenova/phi-3-mini-4k-instruct via Transformers.js",
            examples=[
                "What is the capital of Bangladesh?",
                "Explain AI hallucination simply",
                "Write a Python prime function",
            ],
        )
        
        with gr.Accordion("🔌 API (OpenAI-Compatible)", open=False):
            gr.Markdown(f"""
            **API Base URL:** `https://bdzombie-sahon.hf.space`
            
            - `GET /v1/models` — List models
            - `POST /v1/chat/completions` — Chat
            - `POST /v1/completions` — Text
            - `GET /health` — Health check
            
            All API endpoints served by the **Node.js Transformers.js** backend.
            """)
    
    return demo


# ─── Main Entry Point (wrapped with @spaces.GPU) ───
# The @spaces.GPU decorator keeps GPU allocated for the entire
# lifetime of the server, since uvicorn.run() blocks forever.

def _start_server():
    """Build and start the entire Gradio + FastAPI + Node.js stack."""
    from gradio.routes import App
    import uvicorn
    
    # Build Gradio demo
    demo = run_app()
    
    # Create FastAPI app from Gradio
    fastapi_app = App.create_app(demo)
    
    # ── Proxy: GET /health ──
    @fastapi_app.get("/health")
    async def proxy_health():
        try:
            async with httpx.AsyncClient() as client:
                r = await client.get(f"{NODE_SERVER_URL}/health", timeout=5)
                return JSONResponse(content=r.json())
        except:
            return JSONResponse(
                content={"status": "ok", "node_js": "loading", "progress": "Node.js server starting..."},
                status_code=200,
            )
    
    # ── Proxy: GET /v1/models ──
    @fastapi_app.get("/v1/models")
    async def proxy_models():
        try:
            async with httpx.AsyncClient() as client:
                r = await client.get(f"{NODE_SERVER_URL}/v1/models", timeout=5)
                return JSONResponse(content=r.json())
        except Exception as e:
            return JSONResponse(content={"error": str(e)}, status_code=503)
    
    # ── Proxy: POST /v1/chat/completions ──
    @fastapi_app.post("/v1/chat/completions")
    async def proxy_chat(request: Request):
        try:
            body = await request.json()
            async with httpx.AsyncClient() as client:
                r = await client.post(f"{NODE_SERVER_URL}/v1/chat/completions", json=body, timeout=120)
                return JSONResponse(content=r.json(), status_code=r.status_code)
        except httpx.TimeoutException:
            return JSONResponse(content={"error": "Request timeout"}, status_code=504)
        except Exception as e:
            return JSONResponse(content={"error": str(e)}, status_code=500)
    
    # ── Proxy: POST /v1/completions ──
    @fastapi_app.post("/v1/completions")
    async def proxy_completions(request: Request):
        try:
            body = await request.json()
            async with httpx.AsyncClient() as client:
                r = await client.post(f"{NODE_SERVER_URL}/v1/completions", json=body, timeout=120)
                return JSONResponse(content=r.json(), status_code=r.status_code)
        except Exception as e:
            return JSONResponse(content={"error": str(e)}, status_code=500)
    
    # Start!
    port = int(os.environ.get("PORT", 7860))
    print(f"[Sahon] Starting unified server on 0.0.0.0:{port}")
    uvicorn.run(fastapi_app, host="0.0.0.0", port=port, log_level="info")


# Apply @spaces.GPU if available — this keeps GPU alive while uvicorn runs
if HAS_SPACES:
    main = spaces_gpu(_start_server)
else:
    main = _start_server

if __name__ == "__main__":
    main()