import os
import time
import httpx
import json
import re
from fastapi import FastAPI, Request, HTTPException
from fastapi.responses import StreamingResponse, HTMLResponse, JSONResponse
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel
from typing import List, Optional

app = FastAPI(title="KAI STUDIO ROUTER")
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
    expose_headers=["*"]
)

BACKENDS = {
    "qwen": {
        "url": os.environ.get("KAI_QWEN_URL", "https://your-qwen-space.hf.space"),
        "key": os.environ.get("KEY_QWEN"),
        "models": ["Qwen2.5-7B-Instruct"]
    },
    "mistral": {
        "url": os.environ.get("KAI_MISTRAL_URL", "https://your-mistral-space.hf.space"),
        "key": os.environ.get("KEY_MISTRAL"),
        "models": ["Mistral-7B-v0.3"]
    }
}

class Message(BaseModel):
    role: str
    content: str

class ChatRequest(BaseModel):
    messages: List[Message]
    max_tokens: Optional[int] = 400
    temperature: Optional[float] = 0.7
    stream: bool = True
    model: Optional[str] = None

CODE_KEYWORDS = re.compile(r'\b(def|class|import|function|const|let|var|public|private|return|python|javascript|java|c\+\+|sql|html|css|json|api|docker|git|kod|funkcja|program|skrypt|algorytm)\b', re.I)
MATH_KEYWORDS = re.compile(r'\b(oblicz|policz|równanie|całka|pochodna|macierz|prawdopodobieństwo|logarytm|pierwiastek|ile to|=\?|\d+\s*[\+\-\*\/]\s*\d+)\b', re.I)
CREATIVE_KEYWORDS = re.compile(r'\b(napisz|wiersz|opowiadanie|historia|bajka|rap|tekst|piosenka|kreatywny|stwórz|wymyśl|zrób post|dialog)\b', re.I)

def route_model(messages: List[Message], forced_model: Optional[str] = None) -> str:
    if forced_model and forced_model in BACKENDS:
        return forced_model

    user_text = " ".join([m.content for m in messages if m.role == "user"]).lower()

    if CODE_KEYWORDS.search(user_text) or MATH_KEYWORDS.search(user_text):
        return "qwen"

    if CREATIVE_KEYWORDS.search(user_text):
        return "mistral"

    pl_chars = len(re.findall(r'[ąćęłńóśźż]', user_text))
    if pl_chars > 2:
        return "mistral"

    return "qwen"

# --- 2 oddzielne funkcje: stream i non-stream ---
async def proxy_stream(backend_name: str, request_data: dict):
    backend = BACKENDS[backend_name]
    url = f"{backend['url']}/v1/chat/completions"
    headers = {
        "Authorization": f"Bearer {backend['key']}",
        "Content-Type": "application/json"
    }

    async with httpx.AsyncClient(timeout=180.0) as client:
        async with client.stream("POST", url, json=request_data, headers=headers) as r:
            if r.status_code!= 200:
                err = await r.aread()
                raise HTTPException(r.status_code, f"Backend {backend_name} error: {err.decode()}")

            yield f": x-kai-model-used {backend_name}\n\n"
            yield f": connected\n\n"

            async for chunk in r.aiter_bytes():
                yield chunk

async def proxy_non_stream(backend_name: str, request_data: dict):
    backend = BACKENDS[backend_name]
    url = f"{backend['url']}/v1/chat/completions"
    headers = {
        "Authorization": f"Bearer {backend['key']}",
        "Content-Type": "application/json"
    }

    async with httpx.AsyncClient(timeout=180.0) as client:
        r = await client.post(url, json=request_data, headers=headers)
        if r.status_code!= 200:
            raise HTTPException(r.status_code, f"Backend {backend_name} error: {r.text}")
        data = r.json()
        data["kai_model_used"] = backend_name
        return data

@app.get("/", response_class=HTMLResponse)
async def serve_frontend():
    with open("static/index.html", "r", encoding="utf-8") as f:
        return HTMLResponse(content=f.read())

@app.post("/v1/chat/completions")
async def chat_completions(data: ChatRequest, request: Request):
    chosen_backend = route_model(data.messages, data.model)

    request_data = data.dict()
    request_data.pop("model", None)

    if data.stream:
        generator = proxy_stream(chosen_backend, request_data)
        return StreamingResponse(generator, media_type="text/event-stream")
    else:
        result = await proxy_non_stream(chosen_backend, request_data)
        return JSONResponse(content=result)

@app.get("/health")
async def health():
    return {"status": "ok", "backends": list(BACKENDS.keys())}

app.mount("/static", StaticFiles(directory="static"), name="static")