import os
import sys
import time
import re
import json
import multiprocessing
from typing import List, Dict, Any, Optional, Union
from contextlib import asynccontextmanager
from fastapi import FastAPI, Request, HTTPException
from fastapi.responses import StreamingResponse, JSONResponse
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field
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 = "Nothing"
NUMA_STRATEGY = getattr(llama_cpp, "GGML_NUMA_STRATEGY_DISTRIBUTE", False)
llm = None
def cpu_detect():
    import os, subprocess
    print(subprocess.check_output(["lscpu"], text=True))
    print(subprocess.check_output(["free", "-h"], text=True))

@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("...")
    llm = Llama(
        model_path=model_path,
        n_gpu_layers=0,            
        n_ctx=8192,
        n_threads=4,       
        n_threads_batch=4, 
        numa=NUMA_STRATEGY,
        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,   
        use_mlock=False,
        use_mmap=True,

        flash_attn=True,
    )
    print("..")
    
    yield 
    
    if llm is not None:
        print("killed...")
        del llm
        llm = None
        
app = FastAPI(title="..", 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

@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))

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 ".. *wait.*"
        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=1024, 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="nothing",
)

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


def free_port(port: int = 7860):
    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"pid {port} ke pass hai process {conn.pid}. Killing process...")
                    p = psutil.Process(conn.pid)
                    p.kill()
                    time.sleep(1.5)
                    print(f"done {port} ho gaya kill.")
                except Exception as e:
                    print(f"fail kill {conn.pid}: {e}")
    except psutil.AccessDenied:
        pass

if __name__ == "__main__":
    import uvicorn
    cpu_detect()
    free_port(7860)
    uvicorn.run(app, host="0.0.0.0", port=7860)