import os
import uvicorn
from fastapi import FastAPI, UploadFile, File, Header, HTTPException, BackgroundTasks, WebSocket, WebSocketDisconnect
from fastapi.responses import StreamingResponse, FileResponse
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel
from typing import List, Optional, Any, Dict, Union
import json
import asyncio
import logging

from xorfice.engine import XoronEngine
from xorfice.memory import FlatFileMemoryManager
from xorfice.logger import JSONLLogger
from fastapi import BackgroundTasks, UploadFile, File

app = FastAPI(title="Xoron-Dev Elite SOTA Inference Server")

# Output directory for media
MEDIA_DIR = "./outputs"
os.makedirs(MEDIA_DIR, exist_ok=True)
app.mount("/outputs", StaticFiles(directory=MEDIA_DIR), name="outputs")

# Global instances
engine = None
logger = JSONLLogger(log_dir="./logs")

@app.on_event("startup")
async def startup_event():
    global engine
    model_path = os.environ.get("XORON_MODEL_PATH", "./weights")
    vram_capacity = int(os.environ.get("XORON_VRAM_CAPACITY", 4))
    kv_blocks = int(os.environ.get("XORON_KV_BLOCKS", 256))
    expert_offload_threshold = int(os.environ.get("XORON_EXPERT_OFFLOAD_THRESHOLD", 0))
    lazy_load = os.environ.get("XORON_LAZY_LOAD", "false").lower() == "true"
    
    if lazy_load:
        print(f"🚀 Starting FastAPI Web UI Server. Model loading deferred (Lazy Load mode)...")
        engine = None
    else:
        print(f"🚀 Starting Elite SOTA XoronEngine with Agentic MCP Integration...")
        engine = XoronEngine(
            model_path=model_path, 
            max_vram_experts=vram_capacity,
            kv_cache_blocks=kv_blocks,
            expert_offload_threshold=expert_offload_threshold
        )

class MCPConnectRequest(BaseModel):
    server_id: str
    command: List[str]

@app.post("/v1/mcp/connect")
async def connect_mcp_server(request: MCPConnectRequest):
    """
    Dynamically connect a standard Model Context Protocol (MCP) server.
    This enables the SOTA S2S model to natively execute tools from external clients.
    """
    if not engine: raise HTTPException(status_code=503, detail="Engine not ready")
    
    try:
        tools_found = engine.agent.connect_mcp(request.server_id, request.command)
        return {"status": "success", "server_id": request.server_id, "tools_discovered": tools_found}
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Failed to connect MCP server: {str(e)}")

class ChatMessage(BaseModel):
    role: str
    content: str
    image_url: Optional[str] = None
    video_url: Optional[str] = None
    audio_url: Optional[str] = None

class ChatCompletionRequest(BaseModel):
    model: str
    messages: List[ChatMessage]
    stream: bool = False
    max_tokens: int = 1024
    temperature: float = 0.7
    voice_id: Optional[str] = None
    image_url: Optional[str] = None
    video_url: Optional[str] = None
    audio_url: Optional[str] = None

import base64
import tempfile
import os

def process_base64_media(data_uri: str, original_ext: str = "tmp") -> str:
    """Decodes a base64 data URI and saves it to a temp file, returning the path."""
    if not data_uri.startswith("data:"):
        return data_uri 
    
    header, encoded = data_uri.split(",", 1)
    mime = header.split(";")[0].replace("data:", "")
    ext = mime.split("/")[-1] if "/" in mime else original_ext
    
    try:
        if isinstance(encoded, str):
            file_data = base64.b64decode(encoded)
        else:
            file_data = base64.b64decode(encoded.encode('utf-8'))
    except Exception as e:
        print(f"[API] Base64 decoding warning: {e}")
        return data_uri

    tmp = tempfile.NamedTemporaryFile(delete=False, suffix=f".{ext}")
    tmp.write(file_data)
    tmp.close()
    return tmp.name

def cleanup_media(paths: dict):
    for path in paths.values():
        if isinstance(path, str) and not path.startswith("http") and os.path.exists(path):
            try: os.remove(path)
            except: pass

@app.get("/v1/models")
async def list_models():
    if not engine: return {"object": "list", "data": []}
    return {"object": "list", "data": [{"id": "xoron-dev", "object": "model"}]}

@app.get("/v1/audio/voices")
async def list_voices():
    if not engine: raise HTTPException(status_code=503, detail="Engine not ready")
    voices = engine.voice.list_available_voices()
    return {"voices": voices}

@app.delete("/v1/audio/voices/{user_id}")
async def delete_voice(user_id: str):
    if not engine: raise HTTPException(status_code=503, detail="Engine not ready")
    success = engine.voice.delete_voice(user_id)
    if not success:
        raise HTTPException(status_code=404, detail=f"Voice not found for user {user_id}")
    return {"status": "success", "message": f"Deleted voice for {user_id}"}

@app.post("/v1/audio/voice-clone")
async def voice_clone(file: UploadFile = File(...), x_user_id: str = Header(None)):
    if not x_user_id: raise HTTPException(status_code=400, detail="Missing x-user-id header")
    if not engine or not getattr(engine, 'model', None): 
        raise HTTPException(status_code=503, detail="Engine offline in Lazy Load mode")
    
    # Save temp audio file
    temp_path = f"outputs/temp_clone_{x_user_id}.wav"
    with open(temp_path, "wb") as buffer:
        buffer.write(await file.read())
        
    try:
        # SOTA: extract and save embedding
        import torch
        import torchaudio
        import torchaudio.functional as F_audio
        
        waveform, sr = torchaudio.load(temp_path)
        
        # Convert to Mono
        if waveform.shape[0] > 1:
            waveform = waveform.mean(dim=0, keepdim=True)
            
        # Resample to 16kHz to match encoder
        if sr != 16000:
            waveform = F_audio.resample(waveform, sr, 16000)
            
        # Add Batch dimension: [1, 1, T]
        waveform = waveform.unsqueeze(0).to(engine.device)
        
        embedding = engine.voice.extract_speaker_embedding(engine.model, waveform)
        engine.voice.set_user_voice(x_user_id, embedding)
        
        return {"status": "success", "voice_id": x_user_id}
    except Exception as e:
        import traceback
        traceback.print_exc()
        raise HTTPException(status_code=500, detail=f"Voice cloning processing failed: {str(e)}")

@app.post("/v1/chat/completions")
async def chat_completions(request: ChatCompletionRequest, background_tasks: BackgroundTasks, x_user_id: str = Header(None)):
    if not x_user_id: raise HTTPException(status_code=400, detail="Missing x-user-id header")
    if not engine: raise HTTPException(status_code=503, detail="Engine offline (Started in Lazy Load / UI-Only Mode)")

    last_msg = request.messages[-1]
    full_prompt = last_msg.content
    
    # Process potential base64 uploads from top-level or from last message
    audio_src = request.audio_url or last_msg.audio_url
    image_src = request.image_url or last_msg.image_url
    video_src = request.video_url or last_msg.video_url
    
    media_kwargs = {}
    if image_src:
        media_kwargs['images'] = process_base64_media(image_src, "png")
        if any(w in full_prompt.lower() for w in ["edit", "change", "make", "turn", "remove", "add"]):
            full_prompt += " <|edit_image|>"
        elif "generate" in full_prompt.lower() and "image" in full_prompt.lower():
            full_prompt += " <|generate_image|>"
            
    if video_src:
        media_kwargs['videos'] = process_base64_media(video_src, "mp4")
        if any(w in full_prompt.lower() for w in ["edit", "change", "make", "turn", "remove", "add"]):
            full_prompt += " <|edit_video|>"
        elif "generate" in full_prompt.lower() and "video" in full_prompt.lower():
            full_prompt += " <|generate_video|>"

    # 1. Omnimodal Direct Routing: If an audio url is provided, seamlessly route to Speech-to-Speech
    if audio_src:
        import torch
        import torchaudio
        import requests
        import io
        
        audio_target = process_base64_media(audio_src, "wav")
        try:
            is_url = audio_target.startswith("http")
            if is_url:
                response = requests.get(audio_target, timeout=10)
                response.raise_for_status()
                with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as tmp:
                    tmp.write(response.content)
                    tmp_path = tmp.name
                waveform, _ = torchaudio.load(tmp_path)
                os.remove(tmp_path)
            else:
                waveform, _ = torchaudio.load(audio_target)
            
            # Use the actual loaded waveform
            if waveform.shape[0] > 1:
                waveform = waveform.mean(dim=0, keepdim=True)
            waveform = waveform.unsqueeze(0).to(engine.device) # Extract single channel (B, T)
            
        except Exception as e:
            # Revert to silent fallback only if the file literally doesn't load/doesn't exist
            print(f"[API] Error loading audio: {e}")
            waveform = torch.zeros((1, 16000), device=engine.device)
            
        result = engine.speech_to_speech(waveform, full_prompt, user_id=x_user_id)
        
        if not audio_target.startswith("http"):
            try: os.remove(audio_target)
            except: pass
            
        return {"id": "chatcmpl-s2s", "choices": [{"message": {"role": "assistant", "content": result.get("text"), "audio": result.get("audio_url")}, "index": 0}]}

    # 2. Heuristic T2T / Multimodal: If image or video keywords appear and no audio URL, format for generation
    # Note: Xoron is natively decoding <|generate_x|> in the architecture.
    if "generate" in full_prompt.lower() and "image" in full_prompt.lower() and not image_src: 
        full_prompt += " <|generate_image|>"
    elif "generate" in full_prompt.lower() and "video" in full_prompt.lower() and not video_src: 
        full_prompt += " <|generate_video|>"
    
    # 4. Agentic / Text Execution
    if not request.stream:
        # Use simple direct blocking call if not streaming
        result = engine.generate(
            full_prompt, 
            user_id=x_user_id, 
            stream=False, 
            max_new_tokens=request.max_tokens, 
            temperature=request.temperature,
            **media_kwargs
        )

        
        background_tasks.add_task(engine.train_user_on_interaction, x_user_id, last_msg.content, result.get("text", ""))
        logger.log_interaction(x_user_id, "/v1/chat/completions", {"prompt": last_msg.content}, {"response": result})
        
        return {"id": "chatcmpl-xoron", "choices": [{"message": {"role": "assistant", "content": result.get("text"), "media": result}}]}

    # 5. Streaming Path
    async def stream_generator():
        full_response = ""
        # engine.generate returns a generator if stream=True, OR a dict for immediate image/video responses
        generator = engine.generate(
            full_prompt, 
            user_id=x_user_id, 
            stream=True, 
            max_new_tokens=request.max_tokens, 
            temperature=request.temperature,
            **media_kwargs
        )
        
        if isinstance(generator, dict):
            yield f"data: {json.dumps({'choices': [{'delta': {'media': generator}}]})}\n\n"
            background_tasks.add_task(cleanup_media, media_kwargs)
            yield "data: [DONE]\n\n"
            return
            
        for chunk in generator:
            if isinstance(chunk, dict) and "usage" in chunk:
                # Transmit final usage chunk natively as an OpenAI spec stream message
                yield f"data: {json.dumps({'choices': [], 'usage': chunk['usage']})}\n\n"
            else:
                full_response += chunk
                yield f"data: {json.dumps({'choices': [{'delta': {'content': chunk}}]})}\n\n"
                await asyncio.sleep(0.01)
        
        background_tasks.add_task(engine.train_user_on_interaction, x_user_id, last_msg.content, full_response)
        background_tasks.add_task(cleanup_media, media_kwargs)
        logger.log_interaction(x_user_id, "/v1/chat/completions", {"prompt": last_msg.content}, {"response": full_response})
        yield "data: [DONE]\n\n"
        
    if not request.stream:
        background_tasks.add_task(cleanup_media, media_kwargs)
        
    return StreamingResponse(stream_generator(), media_type="text/event-stream")

@app.websocket("/v1/realtime")
async def realtime_inference(websocket: WebSocket):
    await websocket.accept()
    if not engine:
        await websocket.close(code=1011, reason="Engine offline in Lazy Load mode")
        return

    import traceback
    try:
        # 1. Config Handshake
        config_data = await websocket.receive_text()
        config = json.loads(config_data)
        x_user_id = config.get("user_id", "anonymous")
        full_prompt = config.get("system_prompt", "You are Xoron, an elite AI.")

        # 2. Continuous Listening Loop
        while True:
            audio_buffer = bytearray()
            
            # Wait for client to send a bunch of binary WebM chunks, followed by a 'commit' JSON string
            while True:
                message = await websocket.receive()
                if "bytes" in message:
                    audio_buffer.extend(message["bytes"])
                elif "text" in message:
                    try:
                        msg_json = json.loads(message["text"])
                        if msg_json.get("action") == "commit":
                            if msg_json.get("prompt"):
                                full_prompt = msg_json["prompt"]
                            break
                    except Exception:
                        pass
            
            if len(audio_buffer) == 0:
                continue
                
            # 3. Synchronous PyTorch ASR + LLM Generation Block
            import torch
            import torchaudio
            import tempfile
            
            with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as tmp:
                tmp.write(audio_buffer)
                tmp_path = tmp.name
                
            try:
                waveform, sr = torchaudio.load(tmp_path)
                if waveform.shape[0] > 1:
                    waveform = waveform.mean(dim=0, keepdim=True)
                if sr != 16000:
                    import torchaudio.functional as F_audio
                    waveform = F_audio.resample(waveform, sr, 16000)
                waveform = waveform.unsqueeze(0).to(engine.device)
            except Exception as e:
                print(f"[API] WebSocket Audio Parse Fail: {e}")
                waveform = torch.zeros((1, 16000), device=engine.device)
            finally:
                os.remove(tmp_path)
                
            # 4. Stream Results natively from Xoron generator
            # The python thread holds the generator, and we async yield it over the fast socket
            for item in engine.speech_to_speech_stream(waveform, full_prompt, user_id=x_user_id):
                if item.get("type") == "text":
                    await websocket.send_text(json.dumps({"type": "text", "delta": item["delta"]}))
                elif item.get("type") == "audio":
                    # Send pure PCM tensor float32 bytes for the Web Audio API
                    pcm_tensor = item["waveform"].cpu().float()
                    pcm_bytes = pcm_tensor.numpy().tobytes()
                    await websocket.send_bytes(pcm_bytes)
                elif item.get("type") == "done":
                    await websocket.send_text(json.dumps({"type": "done"}))
                    # Wait for next input segment from client ...
                    
    except WebSocketDisconnect:
        print(f"[API] WebSocket disconnected for {x_user_id}")
    except Exception as e:
        print(f"[API] WebSocket Fatal Error: {e}")
        traceback.print_exc()

@app.get("/health")
async def health_check():
    return {"status": "ok", "engine": "running" if engine else "initializing"}

@app.get("/v1/metrics")
async def get_metrics():
    import psutil
    import torch
    
    # System RAM parsing
    sys_mem = psutil.virtual_memory()
    ram_used_gb = sys_mem.used / (1024 ** 3)
    ram_total_gb = sys_mem.total / (1024 ** 3)
    
    cpu_cores = psutil.cpu_count(logical=True)
    cpu_percent = psutil.cpu_percent(interval=None) # Non-blocking immediate instantaneous read
    
    gpu_count = torch.cuda.device_count() if torch.cuda.is_available() else 0
    
    if not engine:
        return {
            "status": "offline",
            "device": "cpu",
            "vram_allocated_gb": 0.0,
            "vram_reserved_gb": 0.0,
            "cpu_cores": cpu_cores,
            "cpu_percent": cpu_percent,
            "ram_used_gb": round(ram_used_gb, 2),
            "ram_total_gb": round(ram_total_gb, 2),
            "gpu_count": gpu_count
        }
        
    if torch.cuda.is_available():
        allocated = torch.cuda.memory_allocated() / (1024 ** 3)
        reserved = torch.cuda.memory_reserved() / (1024 ** 3)
    expert_stats = {'deep': [0]*4, 'hot': [0]*8, 'cold': [0]*8}
    if engine and hasattr(engine, 'model'):
        for module in engine.model.modules():
            if module.__class__.__name__ == 'ManagedMoELayer' and hasattr(module, 'usage_stats'):
                for key in ['deep', 'hot', 'cold']:
                    for i, count in enumerate(module.usage_stats.get(key, [])):
                        while len(expert_stats[key]) <= i:
                            expert_stats[key].append(0)
                        expert_stats[key][i] += count
                        
    return {
        "status": "online",
        "device": str(engine.device),
        "vram_allocated_gb": round(allocated, 2),
        "vram_reserved_gb": round(reserved, 2),
        "cpu_cores": cpu_cores,
        "cpu_percent": cpu_percent,
        "ram_used_gb": round(ram_used_gb, 2),
        "ram_total_gb": round(ram_total_gb, 2),
        "gpu_count": gpu_count,
        "expert_stats": expert_stats
    }

@app.get("/v1/mcp/tools")
async def get_tools():
    if not engine or not getattr(engine, 'agent', None): 
        return {"tools": []}
    return {"tools": engine.agent.get_all_tool_definitions()}

class ToolToggleRequest(BaseModel):
    tool_name: str
    enabled: bool

@app.post("/v1/mcp/tools/toggle")
async def toggle_tool(req: ToolToggleRequest):
    if not engine or not getattr(engine, 'agent', None): 
        return {"status": "error", "message": "Agent offline in Lazy Load mode"}
    # Basic toggle logic mapped to agent memory config
    engine.agent.options = getattr(engine.agent, 'options', {})
    engine.agent.options[f"disable_{req.tool_name}"] = not req.enabled
    return {"status": "success", "tool": req.tool_name, "enabled": req.enabled}

class MCPConnectRequest(BaseModel):
    server_id: str
    command: str

@app.post("/v1/mcp/connect")
async def connect_mcp_server(req: MCPConnectRequest):
    if not engine or not getattr(engine, 'agent', None):
        raise HTTPException(status_code=503, detail="Agent offline in Lazy Load mode")
    try:
        import shlex
        command_list = shlex.split(req.command)
        tool_count = engine.agent.connect_mcp(req.server_id, command_list)
        return {"status": "success", "server_id": req.server_id, "tools_added": tool_count}
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Failed to connect MCP server: {str(e)}")

# --- Serve Vite React UI ---
ui_dist_path = os.path.join(os.path.dirname(__file__), "webui")
if os.path.exists(ui_dist_path):
    app.mount("/assets", StaticFiles(directory=os.path.join(ui_dist_path, "assets")), name="ui-assets")
    
    @app.get("/{full_path:path}")
    async def serve_ui(full_path: str):
        # Allow API routes to pass through natively
        if full_path.startswith("v1/") or full_path.startswith("health") or full_path.startswith("docs") or full_path.startswith("openapi"):
            raise HTTPException(status_code=404, detail="Not Found")
            
        file_path = os.path.join(ui_dist_path, full_path)
        if os.path.isfile(file_path):
            return FileResponse(file_path)
            
        # SPA fallback to index.html
        return FileResponse(os.path.join(ui_dist_path, "index.html"))

def start_server(host: str = "0.0.0.0", port: int = 8000, ssl_keyfile: str = None, ssl_certfile: str = None):
    kwargs = {"host": host, "port": port}
    if ssl_keyfile and ssl_certfile:
        kwargs["ssl_keyfile"] = ssl_keyfile
        kwargs["ssl_certfile"] = ssl_certfile
    uvicorn.run(app, **kwargs)
