# app.py
import os
import typing as t

import gradio as gr
from openai import OpenAI, AsyncOpenAI

VLLM_BASE_URL = os.getenv("VLLM_BASE_URL")
VLLM_MODEL = os.getenv("VLLM_MODEL", "avibe")
OPENAI_API_KEY = ""


def _build_messages(
    user_message: str,
    history: list[tuple[str, str]] | None,
    system_prompt: str | None,
) -> list[dict[str, str]]:
    """Конвертация истории gr.ChatInterface в формат OpenAI Chat API."""
    messages: list[dict[str, str]] = []
    if system_prompt:
        messages.append({"role": "system", "content": system_prompt})

    if history:
        for user_turn, assistant_turn in history:
            if user_turn:
                messages.append({"role": "user", "content": user_turn})
            if assistant_turn:
                messages.append({"role": "assistant", "content": assistant_turn})

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


def _parse_stop_sequences(stop_sequences_raw: str | None) -> list[str] | None:
    """Парсит стоп-последовательности из строки (по запятым или переводам строк)."""
    if not stop_sequences_raw:
        return None
    raw = stop_sequences_raw.replace("\r", "\n").replace("\n", ",")
    items = [s.strip() for s in raw.split(",")]
    result = [s for s in items if s]
    return result or None


def _create_client(base_url: str | None, api_key: str | None):
    """Синхронный клиент OpenAI-совместимого API (если понадобится)."""
    effective_api_key = api_key or os.environ.get("OPENAI_API_KEY") or "EMPTY_KEY"
    kwargs: dict[str, t.Any] = {"api_key": effective_api_key}
    if base_url:
        kwargs["base_url"] = base_url
    return OpenAI(**kwargs)


def _create_async_client(base_url: str | None, api_key: str | None):
    """Асинхронный клиент OpenAI-совместимого API (vLLM/SGLang/TGI совместим)."""
    effective_api_key = api_key or os.environ.get("OPENAI_API_KEY") or "EMPTY_KEY"
    kwargs: dict[str, t.Any] = {"api_key": effective_api_key}
    if base_url:
        kwargs["base_url"] = base_url
    return AsyncOpenAI(**kwargs)


async def generate_response(
    user_message: str,
    history: list[tuple[str, str]] | None,
    temperature: float,
    top_p: float,
    max_tokens: int | float | None,
    system_prompt: str,
):
    """
    Генерация ответа с потоковой выдачей через OpenAI-совместимый vLLM/SGLang/TGI.

    Это async-генератор: Gradio ChatInterface будет отображать промежуточный текст
    по мере поступления токенов (streaming). См. доки ChatInterface. 
    """
    try:
        client = _create_async_client(base_url=VLLM_BASE_URL, api_key=OPENAI_API_KEY)

        messages = _build_messages(
            user_message=user_message,
            history=history,
            system_prompt=system_prompt,
        )

        # Безопасная обработка max_tokens
        max_new_tokens: int | None = None
        if isinstance(max_tokens, (int, float)) and max_tokens > 0:
            max_new_tokens = int(max_tokens)

        create_kwargs: dict[str, t.Any] = {
            "model": VLLM_MODEL,
            "messages": messages,
            "temperature": float(temperature),
            "top_p": float(top_p),
            "max_tokens": max_new_tokens,
            "stream": True,  # просим потоковую выдачу
        }

        # Для AsyncOpenAI при stream=True возвращается асинхронный поток событий
        stream = await client.chat.completions.create(**create_kwargs)

        accumulated_text = ""
        async for chunk in stream:
            try:
                # стандартный путь для openai>=1.** c Chat Completions
                delta = chunk.choices[0].delta
                content_piece = getattr(delta, "content", None)
                if content_piece:
                    accumulated_text += content_piece
                    # возвращаем на каждом шаге наращиваемую строку
                    yield accumulated_text
            except Exception:
                # запасной путь для несовпадающих реализаций стриминга
                choice0 = chunk.choices[0]
                content_piece = None
                if hasattr(choice0, "delta") and isinstance(choice0.delta, dict):
                    content_piece = choice0.delta.get("content")
                elif hasattr(choice0, "message") and isinstance(choice0.message, dict):
                    content_piece = choice0.message.get("content")
                if content_piece:
                    accumulated_text += content_piece
                    yield accumulated_text

        # Если поток завершился без текста — вернём пустую строку
        if not accumulated_text:
            yield ""

    except Exception as e:
        # Показываем текст ошибки прямо в чате
        yield f"Ошибка: {e}"


def build_demo() -> gr.Blocks:
    """Собираем интерфейс Gradio ChatInterface с дополнительными контролами."""
    with gr.Blocks(title="vLLM Chat (Gradio)") as demo:
        gr.Markdown("### Тестим версию avibe после этапа GRPO")

        gr.ChatInterface(
            fn=generate_response,
            submit_btn="Отправить",
            additional_inputs=[
                gr.Slider(
                    minimum=0.0, maximum=2.0, value=0.7, step=0.01, label="temperature"
                ),
                gr.Slider(
                    minimum=0.0, maximum=1.0, value=0.95, step=0.01, label="top_p"
                ),
                gr.Slider(
                    minimum=1, maximum=8192, value=512, step=1, label="max_tokens"
                ),
                gr.Textbox(
                    label="system prompt",
                    placeholder="Вы можете задать стиль и поведение ассистента...",
                    lines=3,
                ),
            ],
        )

    return demo


def main() -> None:
    """Точка входа. На HF Spaces НЕ указываем порт/хост и не используем share=True."""
    demo = build_demo()
    # demo.queue()  # можно включить, если ожидается высокая конкуррентность
    demo.launch(debug=True)


if __name__ == "__main__":
    main()