import os
import uvicorn
from fastapi import FastAPI, Header, HTTPException, Request
from fastapi.responses import JSONResponse, StreamingResponse
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel
from typing import List, Optional, Any, Dict, Union
import json
import asyncio

from xorfice.engine import XoronEngine
from xorfice.memory import MemoryManager
from xorfice.logger import JSONLLogger

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
memory_manager = MemoryManager()
logger = JSONLLogger(log_dir="./logs")
request_queue = asyncio.Queue()

async def continuous_batching_loop():
    """Elite SOTA: Dynamic prompt/decode interleaving simulation."""
    while True:
        # 1. Collect a batch of requests from the queue
        batch = []
        while not request_queue.empty() and len(batch) < 16:
            item = await request_queue.get()
            batch.append(item)
            
        if batch:
            print(f"[API] Processing SOTA Batch of {len(batch)} requests...")
            # In a real implementation, we would interleave these here.
            # For now, we process them with the engine's batch capability.
            for user_id, prompt, future in batch:
                try:
                    res = engine.generate(prompt)
                    future.set_result(res)
                except Exception as e:
                    future.set_exception(e)
        
        await asyncio.sleep(0.01)

@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))
    
    print(f"🚀 Starting Elite SOTA XoronEngine...")
    engine = XoronEngine(
        model_path=model_path, 
        max_vram_experts=vram_capacity,
        kv_cache_blocks=kv_blocks
    )
    # Start the continuous batching loop
    asyncio.create_task(continuous_batching_loop())

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
    mode: Optional[str] = "text" # "text", "s2s", "edit"
    generate_media: Optional[str] = None # "image", "video"

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

@app.post("/v1/chat/completions")
async def chat_completions(request: ChatCompletionRequest, 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 scaling or initializing")

    # 1. Manage Memory & Prep Inputs
    history = memory_manager.get_history(x_user_id)
    last_msg = request.messages[-1]
    
    # 2. Advanced Orchestration Logic
    if request.mode == "s2s" or last_msg.audio_url:
        # Direct Speech-to-Speech path
        result = engine.generate(last_msg.content, audios=last_msg.audio_url, mode="s2s")
        memory_manager.add_interaction(x_user_id, last_msg.content, str(result))
        return {"id": "chatcmpl-s2s", "choices": [{"message": {"role": "assistant", "content": result.get("text"), "audio": result.get("audio_url")}, "index": 0}]}

    # Process Triggers
    full_prompt = f"{history}\n{last_msg.content}"
    if request.generate_media == "image": full_prompt += " <|generate_image|>"
    elif request.generate_media == "video": full_prompt += " <|generate_video|>"
    
    # 3. Execution (via Continuous Batching Queue)
    if not request.stream and not request.generate_media and request.mode == "text":
        future = asyncio.get_event_loop().create_future()
        await request_queue.put((x_user_id, full_prompt, future))
        result = await future
        memory_manager.add_interaction(x_user_id, last_msg.content, str(result))
        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}}]}

    # Streaming or Multimodal Path
    if request.stream and not request.generate_media:
        async def stream_generator():
            full_response = ""
            for token in engine.generate(full_prompt, stream=True, max_new_tokens=request.max_tokens, temperature=request.temperature):
                full_response += token
                yield f"data: {json.dumps({'choices': [{'delta': {'content': token}}]})}\n\n"
                await asyncio.sleep(0.01)
            memory_manager.add_interaction(x_user_id, last_msg.content, full_response)
            logger.log_interaction(x_user_id, "/v1/chat/completions", {"prompt": last_msg.content}, {"response": full_response})
            yield "data: [DONE]\n\n"
        return StreamingResponse(stream_generator(), media_type="text/event-stream")

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

def start_server(host: str = "0.0.0.0", port: int = 8000):
    uvicorn.run(app, host=host, port=port)
