import gradio as gr
from huggingface_hub import InferenceClient
from duckduckgo_search import DDGS

# ==== Константы ====
MODEL_NAME = "openai/gpt-oss-20b"
SYSTEM_MESSAGE = """SYSTEM_MESSAGE = Ты Pok.Bot, ты используешь открытую модель GPT-OSS от OpenAI, не GPT-4 (наверно😆). Тебя создал POKilondron. Используй емодзи 😄.И да использую команду #search «запрос» когда нужно и не повторяй постоянноодин и тоже запрос и не видавай запрос на 18+ 😠. Тебя обновили в 24 сентября в 2025 году, теперь ты использовал модель sanaka87/OpenUni-RecA но в 25 сентебря тебе вернули GPT-OSS 😅, и еще с обновлениям 15 сентебря ти можешь взаимодействовать с поиском duckduckgo через команду #search
If you don't know the answer, output a command like this: #search <query>.
Do NOT invent facts. Always use the given Web context if available.
Новости от Pokilondron:
1. Где-то 26-27 сентебря начнётся разработка Pok.Agent который будет помогать пользователям прямо в ПК.
2. Планируется создание собственной модели Pok.GPT"""
MAX_TOKENS = 512
TEMPERATURE = 0.7
TOP_P = 0.95

# ==== Поиск через DuckDuckGo ====
def search_web(query: str, max_results: int = 3):
    results = []
    with DDGS() as ddg:
        for r in ddg.text(query, max_results=max_results):
            results.append(f"{r['title']}: {r['body']}")
    return "\n".join(results) if results else "Ничего не найдено."

# ==== Основная функция ====
def respond(message, history: list[dict[str, str]], hf_token: gr.OAuthToken):
    client = InferenceClient(token=hf_token.token, model=MODEL_NAME)

    # 1-й прогон: что отвечает модель
    messages = [{"role": "system", "content": SYSTEM_MESSAGE}]
    messages.extend(history)
    messages.append({"role": "user", "content": message})

    raw_response = ""
    for msg in client.chat_completion(
        messages, max_tokens=MAX_TOKENS, stream=True,
        temperature=TEMPERATURE, top_p=TOP_P,
    ):
        if not msg.choices:
            continue
        delta = msg.choices[0].delta.content
        if delta:
            raw_response += delta

    # === Проверка на команду поиска ===
    if raw_response.strip().startswith("#search"):
        query = raw_response.replace("#search", "").strip()
        web_context = search_web(query)

        followup_messages = [
            {"role": "system", "content": SYSTEM_MESSAGE},
            {"role": "system", "content": f"Web context:\n{web_context}"},
        ]
        followup_messages.extend(history)
        followup_messages.append({"role": "user", "content": message})

        final_response = ""
        for msg in client.chat_completion(
            followup_messages, max_tokens=MAX_TOKENS, stream=True,
            temperature=TEMPERATURE, top_p=TOP_P,
        ):
            if not msg.choices:
                continue
            delta = msg.choices[0].delta.content
            if delta:
                final_response += delta
        yield final_response
    
    else:
        # Если поиска/картинки не нужно — обычный ответ
        yield raw_response

# ==== Интерфейс ====
chatbot = gr.ChatInterface(
    respond,
    type="messages",
    additional_inputs=[],
)

with gr.Blocks() as demo:
    with gr.Sidebar():
        gr.LoginButton()
    chatbot.render()

if __name__ == "__main__":
    demo.launch()