#!/usr/bin/env python3
"""
server.py — High-performance FastAPI backend & Modern Web UI for Mac Local Inference.
Provides:
  - REST API endpoint (/api/chat)
  - Real-time token streaming endpoint (/api/chat/stream) via Server-Sent Events
  - System status endpoint (/api/status)
  - Built-in Swagger docs (/docs)
  - Serves modern, glassmorphic Web UI (/)
"""

import os
import sys
import json
import time
import base64
import asyncio
import threading
from typing import List, Optional, Dict, Any
from pathlib import Path

from fastapi import FastAPI, Request, HTTPException
from fastapi.responses import HTMLResponse, StreamingResponse, JSONResponse, FileResponse
from fastapi.staticfiles import StaticFiles
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel

# Import MLX loader and generator from chat module
from chat import SYSTEM_PROMPT, try_mlx, try_transformers

# Voice-to-Voice engine (Whisper STT + Piper TTS)
from voice import engine as voice_engine

# Instruct the brain to answer in a spoken, TTS-friendly way during voice mode
VOICE_INSTRUCTION = (
    "You are speaking to the user aloud by voice right now. "
    "Answer as a natural spoken conversation: be concise, friendly, and clear. "
    "Use short, simple sentences. Plain text only — NO markdown, code blocks, "
    "headings, bullet points, emojis, or special symbols. "
    "IMPORTANT: Reply in the SAME language the user just spoke in "
    "(e.g. if they speak Hindi, answer in Hindi)."
)

# Global model state
backend = None
model_info = {
    "name": "Loading...",
    "backend": "none",
    "path": "",
    "status": "initializing"
}

BASE_DIR = Path(__file__).resolve().parent
STATIC_DIR = BASE_DIR / "static"
MODELS_DIR = BASE_DIR / "models" / "swastik-0.5b"

def load_active_model():
    """Load the best available local model."""
    global backend, model_info
    
    target_path = None
    model_name = "Flux91/swastik-0.5b"
    
    if MODELS_DIR.is_dir() and (MODELS_DIR / "model.safetensors").exists():
        target_path = str(MODELS_DIR)
        model_name = "Flux91/swastik-0.5b (Fine-Tuned MLX 4-bit)"
    else:
        # Fallback to remote or default
        target_path = "Flux91/swastik-0.5b"
        model_name = "Flux91/swastik-0.5b"

    print(f"🔄 Initializing model from: {target_path}")
    b, _ = try_mlx(target_path)
    if b is None:
        print("⚠️ MLX failed, attempting transformers fallback...")
        b, _ = try_transformers(target_path)

    if b is not None:
        backend = b
        model_info = {
            "name": model_name,
            "backend": b[0].upper(),
            "path": target_path,
            "status": "ready"
        }
        print(f"✅ Model loaded successfully via {model_info['backend']}!")
    else:
        model_info["status"] = "error"
        print("❌ Could not load model.")

# Initialize model at startup
load_active_model()

app = FastAPI(
    title="AI Chatbot API",
    description="High-performance local inference powered by Apple Silicon MLX",
    version="1.0.0"
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Mount static folder
STATIC_DIR.mkdir(parents=True, exist_ok=True)
app.mount("/static", StaticFiles(directory=str(STATIC_DIR)), name="static")

class MessageItem(BaseModel):
    role: str
    content: Any

class ChatRequest(BaseModel):
    message: str
    history: Optional[List[MessageItem]] = []
    temperature: Optional[float] = 0.7
    max_tokens: Optional[int] = 512
    system_prompt: Optional[str] = None

def build_sanitized_messages(message: str, history: List[MessageItem], system_prompt: Optional[str] = None):
    """Build a sanitized list of {role, content: str} messages. No system prompt injected by default."""
    messages = []
    
    # Only add system prompt if user explicitly provided one in settings
    if system_prompt and system_prompt.strip():
        messages.append({"role": "system", "content": system_prompt.strip()})
    
    if history:
        for item in history:
            role = item.role
            c = item.content
            if isinstance(c, list):
                c = "".join(p.get("text", "") if isinstance(p, dict) else str(p) for p in c)
            elif isinstance(c, dict):
                c = c.get("text", str(c))
            messages.append({"role": role, "content": str(c)})

    user_c = message
    if isinstance(user_c, list):
        user_c = "".join(p.get("text", "") if isinstance(p, dict) else str(p) for p in user_c)
    elif isinstance(user_c, dict):
        user_c = user_c.get("text", str(user_c))
    messages.append({"role": "user", "content": str(user_c)})
    
    return messages

@app.get("/", response_class=HTMLResponse)
async def index():
    """Serve the main single page web application."""
    index_file = STATIC_DIR / "index.html"
    if index_file.exists():
        return FileResponse(str(index_file))
    return HTMLResponse("<h1>AI Chatbot API is running. Missing index.html in static/</h1>")

@app.get("/api/status")
async def get_status():
    """Return current server and model status."""
    return {
        "status": model_info["status"],
        "model_name": model_info["name"],
        "backend": model_info["backend"],
        "path": model_info["path"],
        "apple_silicon": True
    }

@app.post("/api/chat")
async def chat_endpoint(req: ChatRequest):
    """Standard JSON API chat endpoint."""
    global backend
    if backend is None:
        raise HTTPException(status_code=503, detail="Model is not loaded.")
    
    backend_type = backend[0]
    messages = build_sanitized_messages(req.message, req.history, req.system_prompt)
    
    start_t = time.time()
    if backend_type == "mlx":
        from chat import generate_mlx
        reply = generate_mlx(backend, messages, req.max_tokens, req.temperature)
    else:
        from chat import generate_transformers
        reply = generate_transformers(backend, messages, req.max_tokens, req.temperature)
    duration = time.time() - start_t
    
    return {
        "reply": reply,
        "model": model_info["name"],
        "backend": model_info["backend"],
        "latency_sec": round(duration, 3)
    }

@app.post("/api/chat/stream")
async def chat_stream_endpoint(req: ChatRequest):
    """Real-time token streaming endpoint using Server-Sent Events (SSE)."""
    global backend
    if backend is None:
        raise HTTPException(status_code=503, detail="Model is not loaded.")
    
    backend_type = backend[0]
    messages = build_sanitized_messages(req.message, req.history, req.system_prompt)
    
    if backend_type != "mlx":
        # Real-time token streaming for PyTorch / Transformers
        from chat import stream_transformers

        async def trans_stream():
            loop = asyncio.get_event_loop()
            token_queue = asyncio.Queue()

            def run_generation():
                try:
                    for tok in stream_transformers(backend, messages, req.max_tokens, req.temperature):
                        loop.call_soon_threadsafe(token_queue.put_nowait, tok)
                except Exception as e:
                    print(f"⚠️ Streaming error in transformers backend: {e}")
                finally:
                    loop.call_soon_threadsafe(token_queue.put_nowait, None)

            threading.Thread(target=run_generation, daemon=True).start()

            while True:
                tok = await token_queue.get()
                if tok is None:
                    break
                payload = json.dumps({"token": tok})
                yield f"data: {payload}\n\n"

            yield "data: [DONE]\n\n"

        return StreamingResponse(
            trans_stream(),
            media_type="text/event-stream",
            headers={
                "Cache-Control": "no-cache",
                "Connection": "keep-alive",
                "X-Accel-Buffering": "no"
            }
        )

    # Real MLX token streaming
    _, model, tokenizer = backend
    if hasattr(tokenizer, 'apply_chat_template'):
        prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    else:
        prompt = "".join(f"<|im_start|>{m['role']}\n{m['content']}<|im_end|>\n" for m in messages) + "<|im_start|>assistant\n"

    from mlx_lm.generate import stream_generate
    from mlx_lm.sample_utils import make_sampler
    sampler = make_sampler(temp=req.temperature)

    async def token_generator():
        # Run generator in separate thread to prevent blocking the async event loop
        loop = asyncio.get_event_loop()
        def generate_all_tokens():
            tokens = []
            for chunk in stream_generate(model, tokenizer, prompt, max_tokens=req.max_tokens, sampler=sampler):
                tokens.append(chunk.text)
            return tokens

        # Yield tokens as generated
        for chunk in stream_generate(model, tokenizer, prompt, max_tokens=req.max_tokens, sampler=sampler):
            payload = json.dumps({"token": chunk.text})
            yield f"data: {payload}\n\n"
            await asyncio.sleep(0.005) # Yield control to event loop
        
        yield "data: [DONE]\n\n"

    return StreamingResponse(
        token_generator(),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "Connection": "keep-alive",
            "X-Accel-Buffering": "no"
        }
    )

# ════════════════════════════════════════════════════════════
#  VOICE-TO-VOICE ENDPOINTS
# ════════════════════════════════════════════════════════════

@app.get("/api/voice/status")
async def voice_status():
    """Voice components status + download progress."""
    s = voice_engine.status()
    s["piper_available"] = voice_engine._piper is not None
    s["say_available"] = voice_engine._say_ok
    return s


@app.post("/api/voice/init")
async def voice_init():
    """Start downloading/loading voice models (idempotent, returns immediately)."""
    voice_engine.start_init()
    return voice_engine.status()


class VoiceRoundRequest(BaseModel):
    audio: str                              # base64 WAV (16-bit PCM, 16 kHz mono)
    history: Optional[List[MessageItem]] = []
    system_prompt: Optional[str] = None
    temperature: Optional[float] = 0.7
    max_tokens: Optional[int] = 256


@app.post("/api/voice/round")
async def voice_round(req: VoiceRoundRequest, request: Request):
    """One voice round: STT -> LLM -> TTS, streamed as Server-Sent Events.

    Each `data:` line is a JSON object with a `type` field:
      loading  → "loading" events during model download/loading
      sts      → {text, language} Whisper English transcript
      ai_start → LLM generation started
      token    → {text} live token for the overlay transcript
      audio    → {data} base64 WAV (sentence-sized) spoken via TTS
      ai_end   → {text} full assistant reply
      done     → stream complete
      error    → {message}
    """
    global backend

    try:
        wav_bytes = base64.b64decode(req.audio)
    except Exception:
        raise HTTPException(status_code=422, detail="Invalid base64 audio.")

    if backend is None:
        raise HTTPException(status_code=503, detail="LLM not loaded.")

    LANG_NAMES = {
        "en": "English", "hi": "Hindi", "ta": "Tamil", "bn": "Bengali",
        "mr": "Marathi", "te": "Telugu", "kn": "Kannada", "gu": "Gujarati",
        "ur": "Urdu", "pa": "Punjabi", "ml": "Malayalam", "or": "Odia",
        "as": "Assamese", "ne": "Nepali", "hi-en": "Hindi",
    }

    def build_voice_messages(user_text, user_history, lang_code="en"):
        lang_name = LANG_NAMES.get((lang_code or "en").lower().split("-")[0])
        msgs = []
        sys_parts = []
        if req.system_prompt and req.system_prompt.strip():
            sys_parts.append(req.system_prompt.strip())
        sys_parts.append(VOICE_INSTRUCTION)
        if lang_name and lang_name != "English":
            sys_parts.append(
                f"The user just spoke to you in {lang_name} (language code '{lang_code}'). "
                f"Your ENTIRE spoken reply must be in {lang_name} — same language and script "
                f"as the user. Do NOT switch to English. Answer naturally, in a sentence or two."
            )
        msgs.append({"role": "system", "content": "\n\n".join(sys_parts)})
        for item in user_history or []:
            c = item.content
            if isinstance(c, list):
                c = "".join(p.get("text", "") if isinstance(p, dict) else str(p) for p in c)
            elif isinstance(c, dict):
                c = c.get("text", str(c))
            msgs.append({"role": item.role, "content": str(c)})
        msgs.append({"role": "user", "content": user_text})
        return msgs

    async def event_stream():
        def sse(payload):
            return f"data: {json.dumps(payload)}\n\n"

        loop = asyncio.get_event_loop()
        queue = asyncio.Queue()

        def worker():
            def put(obj):
                loop.call_soon_threadsafe(queue.put_nowait, obj)
            try:
                # 1. Ensure voice models loaded
                if not voice_engine.is_ready():
                    voice_engine.ensure_loaded(progress_cb=lambda st: put({"type": "loading", **st}))
                if not voice_engine.is_ready():
                    put({"type": "error", "message": "Voice models could not be loaded."})
                    return

                # 2. STT (Whisper — keeps the user's own language)
                try:
                    stt = voice_engine.transcribe(wav_bytes)
                except Exception as e:
                    put({"type": "error", "message": "Transcription failed: " + str(e)})
                    return
                user_text = (stt.get("text") or "").strip()
                language = stt.get("language") or "en"
                if not user_text:
                    put({"type": "error", "message": "Could not understand the audio."})
                    return
                put({"type": "sts", "text": user_text, "language": language})

                # Ensure the TTS voice for the detected language is available
                try:
                    voice_engine.ensure_lang(language, lambda st: put({"type": "loading", **st}))
                except Exception as e:
                    print(f"  [voice] could not load voice for language '{language}': {e}")

                # 3. LLM generation + per-sentence TTS
                put({"type": "ai_start"})
                messages = build_voice_messages(user_text, req.history or [], language)
                b_type = backend[0]

                if b_type == "mlx":
                    from mlx_lm.sample_utils import make_sampler
                    from mlx_lm.generate import stream_generate
                    _, model, tokenizer = backend
                    prompt = tokenizer.apply_chat_template(
                        messages, tokenize=False, add_generation_prompt=True
                    )
                    sampler = make_sampler(temp=req.temperature)
                    full = ""
                    sentence = ""
                    for chunk in stream_generate(
                        model, tokenizer, prompt,
                        max_tokens=req.max_tokens, sampler=sampler,
                    ):
                        token = chunk.text
                        full += token
                        sentence += token
                        put({"type": "token", "text": token})
                        if sentence.strip() and (is_sentence_end(sentence) or len(sentence) >= 180):
                            for wav_b64 in voice_engine.synth_chunks(sentence, language):
                                put({"type": "audio", "data": wav_b64, "language": language})
                            sentence = ""
                    if sentence.strip():
                        for wav_b64 in voice_engine.synth_chunks(sentence, language):
                            put({"type": "audio", "data": wav_b64, "language": language})
                    put({"type": "ai_end", "text": full})
                    put({"type": "done"})
                else:
                    put({"type": "error", "message": "Voice mode requires the MLX backend."})

            except Exception as e:
                import traceback
                traceback.print_exc()
                put({"type": "error", "message": str(e)})

        threading.Thread(target=worker, daemon=True).start()

        while True:
            if await request.is_disconnected():
                break
            item = await queue.get()
            yield sse(item)
            if item.get("type") in ("done", "error"):
                break

    return StreamingResponse(
        event_stream(),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "Connection": "keep-alive",
            "X-Accel-Buffering": "no"
        }
    )


def is_sentence_end(text):
    """True if text ends with a sentence-terminating punctuation."""
    return text.rstrip().endswith((".", "!", "?", "।", "。", "！", "？"))


if __name__ == "__main__":
    import uvicorn
    default_port = 7860 if os.getenv("SPACE_ID") or os.getenv("HF_SPACE") else 8000
    port = int(os.getenv("PORT", default_port))
    print(f"\n🚀 AI Chatbot Server running at http://localhost:{port}")
    print(f"📖 Interactive API Documentation at http://localhost:{port}/docs\n")
    uvicorn.run(app, host="0.0.0.0", port=port, log_level="info")