import os
import sys
import time
import re
import json
import multiprocessing
from typing import List, Dict, Any, Optional, Union
from contextlib import asynccontextmanager

# FastAPI & Pydantic imports
from fastapi import FastAPI, Request, HTTPException
from fastapi.responses import StreamingResponse, JSONResponse
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field

# ---------------------------------------------------------------------------
# ⚡ 1. CPU & THREADING OPTIMIZATIONS FOR MoE
# ---------------------------------------------------------------------------
try:
    import psutil
    PHYSICAL_CORES = psutil.cpu_count(logical=False) or max(1, multiprocessing.cpu_count() // 2)
except Exception:
    PHYSICAL_CORES = max(1, multiprocessing.cpu_count() // 2)

os.environ["OMP_NUM_THREADS"] = str(PHYSICAL_CORES)
os.environ["MKL_NUM_THREADS"] = str(PHYSICAL_CORES)
os.environ["OPENBLAS_NUM_THREADS"] = str(PHYSICAL_CORES)
os.environ["OMP_PROC_BIND"] = "TRUE"

try:
    p = psutil.Process()
    if sys.platform == "win32":
        p.nice(psutil.HIGH_PRIORITY_CLASS)
    else:
        p.nice(-15)
except Exception:
    pass

import gradio as gr
import llama_cpp
from llama_cpp import Llama
from huggingface_hub import hf_hub_download

MODEL_NAME = "arush/jarush"
NUMA_STRATEGY = getattr(llama_cpp, "GGML_NUMA_STRATEGY_DISTRIBUTE", False)

# Initialize the global LLM variable
llm = None
def cpu_detect():
    import os, subprocess
    print(subprocess.check_output(["lscpu"], text=True))
    print(subprocess.check_output(["free", "-h"], text=True))


# ---------------------------------------------------------------------------
# ⚡ 2. FASTAPI LIFESPAN (MODEL DOWNLOAD & INITIALIZATION)
# ---------------------------------------------------------------------------
@asynccontextmanager
async def lifespan(app: FastAPI):
    global llm
    
    print("Downloading MoE model GGUF...")
    model_path = hf_hub_download(
        repo_id='mradermacher/Qwen3.6-35B-A3B-Uncensored-HauhauCS-FP16-GGUF',
        filename="Qwen3.6-35B-A3B-Uncensored-HauhauCS-FP16.Q3_K_S.gguf",
    )
    print("Model downloaded successfully!")

    print("Initializing llama-cpp MoE Model...")
    llm = Llama(
        model_path=model_path,
        n_gpu_layers=0,             # 0 for CPU execution
        n_ctx=8192,
        
        # Threading optimizations
        n_threads=4,       
        n_threads_batch=4, 
        numa=NUMA_STRATEGY,
        
        # ⚡ TURBOQUANT EQUIVALENT SETTINGS:
        # Quantizing the KV cache to Q4_0 shrinks memory overhead and drastically
        # reduces RAM bandwidth usage, supercharging CPU generation speeds.
        type_k=getattr(llama_cpp, "GGML_TYPE_Q8_0", 2),
        type_v=getattr(llama_cpp, "GGML_TYPE_Q8_0", 2),
        
        n_batch=512,    
        n_ubatch=512,   
        #n_draft=3,
        
        # ⚡ A3B MoE OPTIMIZATION: 
        # use_mmap=True ensures the OS only pages the active 3B parameters 
        # from disk to RAM, keeping CPU cache thrashing to an absolute minimum.
        use_mlock=False,
        use_mmap=True,
        
        # FlashAttention removed as requested
        flash_attn=True,
    )
    print("Server ready!")
    
    yield 
    
    if llm is not None:
        print("Cleaning up Llama model memory...")
        del llm
        llm = None

# ---------------------------------------------------------------------------
# ⚡ 3. FASTAPI INITIALIZATION & OPENAI PYDANTIC SCHEMAS
# ---------------------------------------------------------------------------
app = FastAPI(title="OpenAI-Compatible CPU MoE Server", version="1.0.0", lifespan=lifespan)

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

class ChatMessage(BaseModel):
    role: str
    content: Union[str, List[Dict[str, Any]]]

class ChatCompletionRequest(BaseModel):
    model: Optional[str] = MODEL_NAME
    messages: List[ChatMessage]
    temperature: Optional[float] = 0.7
    top_p: Optional[float] = 0.95
    max_tokens: Optional[int] = 512
    stream: Optional[bool] = False

# ---------------------------------------------------------------------------
# ⚡ 4. OPENAI-COMPATIBLE REST ENDPOINTS
# ---------------------------------------------------------------------------
@app.get("/v1/models")
async def list_models():
    return {
        "object": "list",
        "data": [{"id": MODEL_NAME, "object": "model", "created": int(time.time()), "owned_by": "local"}],
    }

async def openai_stream_generator(messages_dict: list, max_tokens: int, temperature: float, top_p: float):
    global llm
    if llm is None:
        yield f"data: {json.dumps({'error': 'Model loading in progress'})}\n\n"
        return

    for chunk in llm.create_chat_completion(
        messages=messages_dict,
        max_tokens=max_tokens,
        temperature=temperature,
        top_p=top_p,
        stream=True,
    ):
        yield f"data: {json.dumps(chunk)}\n\n"
    yield "data: [DONE]\n\n"

@app.post("/v1/chat/completions")
async def chat_completions(req: ChatCompletionRequest):
    global llm
    if llm is None:
        raise HTTPException(status_code=503, detail="Model is currently loading or unavailable.")

    try:
        messages_dict = [{"role": m.role, "content": m.content} for m in req.messages]
        
        if req.stream:
            return StreamingResponse(
                openai_stream_generator(messages_dict, req.max_tokens, req.temperature, req.top_p),
                media_type="text/event-stream",
            )
        else:
            response = llm.create_chat_completion(
                messages=messages_dict,
                max_tokens=req.max_tokens,
                temperature=req.temperature,
                top_p=req.top_p,
                stream=False,
            )
            return JSONResponse(content=response)
            
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

# ---------------------------------------------------------------------------
# ⚡ 5. GRADIO UI & HISTORY CLEANER
# ---------------------------------------------------------------------------
def clean_assistant_reply(text):
    if not text:
        return ""
    if isinstance(text, list):
        extracted = []
        for part in text:
            if isinstance(part, str):
                extracted.append(part)
            elif isinstance(part, dict):
                extracted.append(part.get("text", part.get("content", "")))
        text = "".join(extracted)
    elif not isinstance(text, str):
        text = str(text)

    cleaned = re.sub(r'<think>.*?</think>\n*', '', text, flags=re.DOTALL)
    cleaned = cleaned.replace("</think>", "")
    cleaned = re.sub(r'\n\n---\n\*⏱️.*?\*$', '', cleaned)
    return cleaned.strip()

def respond(message, history, system_message, max_tokens, temperature, top_p):
    global llm
    if llm is None:
        yield "⚠️ *The model is currently loading. Please wait a few seconds and try again.*"
        return

    messages = [{"role": "system", "content": system_message}]
    
    for item in history:
        if isinstance(item, dict):
            role = item.get("role", "user")
            content = item.get("content", "")
            if role == "assistant":
                content = clean_assistant_reply(content)
            messages.append({"role": role, "content": content})
        elif isinstance(item, (list, tuple)):
            if item[0]:
                messages.append({"role": "user", "content": item[0]})
            if item[1]:
                messages.append({"role": "assistant", "content": clean_assistant_reply(item[1])})

    messages.append({"role": "user", "content": message})

    response = ""
    token_count = 0
    start_time = time.time()
    
    for chunk in llm.create_chat_completion(
        messages=messages,
        max_tokens=max_tokens,
        temperature=temperature,
        top_p=top_p,
        stream=True,
    ):
        delta = chunk["choices"][0]["delta"]
        if "content" in delta and delta["content"]:
            response += delta["content"]
            token_count += 1
            yield response
            
    end_time = time.time()
    elapsed_time = end_time - start_time
    
    if elapsed_time > 0 and token_count > 0:
        tps = token_count / elapsed_time
        response += f"\n\n---\n*⏱️ {tps:.2f} tok/s ({token_count} tokens in {elapsed_time:.2f}s)*"
        yield response

chatbot = gr.ChatInterface(
    fn=respond,
    additional_inputs=[
        gr.Textbox(value="You are a helpful coding assistant.", label="System message"),
        gr.Slider(minimum=1, maximum=8192, value=512, step=1, label="Max new tokens"),
        gr.Slider(minimum=0.1, maximum=2.0, value=0.7, step=0.1, label="Temperature"),
        gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p"),
    ],
    title="FREEAPI RUNNING ON CPU Qwen 35B-A3B",
)

app = gr.mount_gradio_app(app, chatbot, path="/")

# ---------------------------------------------------------------------------
# ⚡ 6. SERVER LAUNCH (WITH AUTOMATIC PORT CLEANUP)
# ---------------------------------------------------------------------------
def free_port(port: int = 7860):
    """Kills any lingering background process holding the specified port."""
    current_pid = os.getpid()
    try:
        for conn in psutil.net_connections(kind='inet'):
            if conn.laddr and conn.laddr.port == port and conn.pid and conn.pid != current_pid:
                try:
                    print(f"⚠️ Port {port} is held by PID {conn.pid}. Killing process...")
                    p = psutil.Process(conn.pid)
                    p.kill()
                    time.sleep(1.5)  # Wait for OS to release the socket
                    print(f"✅ Port {port} freed successfully.")
                except Exception as e:
                    print(f"Could not kill process {conn.pid}: {e}")
    except psutil.AccessDenied:
        pass # Ignore if run in restricted environment without root

if __name__ == "__main__":
    import uvicorn
    cpu_detect()
    # Clean up port 7860 before trying to launch
    free_port(7860)
    
    # Port 7860 is required by Hugging Face Spaces for health checks
    uvicorn.run(app, host="0.0.0.0", port=7860)