from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
import os
from dotenv import load_dotenv
import requests
from typing import Dict, Any, List
from pydantic import BaseModel
import time
import json

load_dotenv()

app = FastAPI()

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

# Получаем переменные окружения
FLOWISE_API_BASE_URL = os.getenv("FLOWISE_API_BASE_URL")
FLOWISE_CHATFLOW_ID = os.getenv("FLOWISE_CHATFLOW_ID")

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

class ChatCompletionRequest(BaseModel):
    model: str
    messages: List[ChatMessage]
    temperature: float = 0.7

def count_tokens(text: str) -> int:
    # Простой подсчет токенов (слова + знаки препинания)
    return len(text.split()) + len([c for c in text if c in ".,!?;:()[]{}"])

def clean_assistant_response(text: str) -> str:
    # Удаляем лишние маркеры кода и форматирования
    text = text.strip()
    if text.endswith("```"):
        text = text[:-3].strip()
    return text

class CustomJSONResponse(JSONResponse):
    media_type = "application/json; charset=utf-8"

    def render(self, content: Any) -> bytes:
        return json.dumps(
            content,
            ensure_ascii=False,
            allow_nan=False,
            indent=None,
            separators=(',', ':')
        ).encode('utf-8')

@app.get("/")
async def root():
    return CustomJSONResponse({"status": "FastFlowWrapper is running"})

@app.get("/v1/models")
async def get_models():
    try:
        # Запрашиваем список чатфлоу из Flowise
        response = requests.get(f"{FLOWISE_API_BASE_URL}/chatflows")
        response.raise_for_status()
        chatflows = response.json()
        
        # Преобразуем в формат OpenAI API
        models = []
        for chatflow in chatflows:
            models.append({
                "id": chatflow.get("id"),
                "object": "model",
                "created": int(time.time()),
                "owned_by": "flowise",
                "permission": [],
                "root": "flowise",
                "parent": None,
                "system_fingerprint": "phi4-r1"
            })
        
        return CustomJSONResponse({"object": "list", "data": models})
    except requests.RequestException as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/v1/chat/completions")
async def create_chat_completion(request: ChatCompletionRequest):
    try:
        # Получаем последнее сообщение из диалога
        last_message = request.messages[-1]
        if last_message.role != "user":
            raise HTTPException(status_code=400, detail="Last message must be from user")
        
        # Подсчитываем токены запроса
        prompt_tokens = count_tokens(last_message.content)
        
        # Формируем запрос к Flowise
        flowise_request = {
            "question": last_message.content
        }
        
        # Засекаем время начала запроса
        start_time = time.time()
        
        # Отправляем запрос к Flowise с таймаутом
        response = requests.post(
            f"{FLOWISE_API_BASE_URL}/prediction/{FLOWISE_CHATFLOW_ID}",
            json=flowise_request,
            timeout=10  # Уменьшаем таймаут до 10 секунд
        )
        response.raise_for_status()
        
        # Получаем и очищаем ответ
        flowise_response = response.json()
        assistant_response = clean_assistant_response(flowise_response.get("text", ""))
        
        # Подсчитываем токены ответа
        completion_tokens = count_tokens(assistant_response)
        
        return CustomJSONResponse({
            "id": "chatcmpl-" + os.urandom(12).hex(),
            "object": "chat.completion",
            "created": int(start_time),
            "model": "phi4-r1",
            "choices": [
                {
                    "index": 0,
                    "logprobs": None,
                    "finish_reason": "stop",
                    "message": {
                        "role": "assistant",
                        "content": assistant_response
                    }
                }
            ],
            "usage": {
                "prompt_tokens": prompt_tokens,
                "completion_tokens": completion_tokens,
                "total_tokens": prompt_tokens + completion_tokens
            },
            "stats": {},
            "system_fingerprint": "phi4-r1"
        })
    except requests.RequestException as e:
        raise HTTPException(status_code=500, detail=str(e))