import spaces
import json
import subprocess
from llama_cpp import Llama
from llama_cpp_agent import LlamaCppAgent
from llama_cpp_agent import MessagesFormatterType
from llama_cpp_agent.providers import LlamaCppPythonProvider
from llama_cpp_agent.chat_history import BasicChatHistory
from llama_cpp_agent.chat_history.messages import Roles
import gradio as gr
from huggingface_hub import hf_hub_download

# UI элементтері үшін қазақша мәтін
PLACEHOLDER_KZ = """
<div style='padding: 20px; border: 1px solid #e0e0e0; border-radius: 10px; background-color: #f9f9f9; text-align: center;'>
    <h2 style='color: #333;'>🐬 Dolphin-мен сұхбат</h2>
    <p style='color: #555; font-size: 1.1em;'>Төмендегіге хабарламаңызды теріңіз...</p>
</div>
"""

llm = None
llm_model = None

# Модельдерді жүктеу (бұл процесс CPU-да ұзақ уақыт алуы мүмкін)
hf_hub_download(repo_id="bartowski/cognitivecomputations_Dolphin3.0-R1-Mistral-24B-GGUF", filename="cognitivecomputations_Dolphin3.0-R1-Mistral-24B-Q8_0.gguf",  local_dir = "./models")
hf_hub_download(repo_id="cognitivecomputations/dolphin-2.9.2-qwen2-72b-gguf", filename="qwen2-Q3_K_M.gguf",  local_dir = "./models")

# @spaces.GPU(duration=120) # <-- Бұл жолды алып тастадық (zero-GPU үшін)
def respond(
    message,
    history: list[tuple[str, str]],
    model,
    max_tokens,
    temperature,
    top_p,
    top_k,
    repeat_penalty,
):
    global llm
    global llm_model

    if llm is None or llm_model != model:
        llm = Llama(
            model_path=f"models/{model}",
            flash_attn=False, # <-- CPU үшін 'False'
            n_gpu_layers=0,  # <-- CPU үшін '0' (GPU қабаты жоқ)
            n_batch=1024,
            n_ctx=8192,
        )
        llm_model=model
    provider = LlamaCppPythonProvider(llm)

    # System Prompt қазақшаға аударылды
    system_prompt_kz = """Сіз Dolphinсіз, адамзатқа көмектесетін жасанды интеллект көмекшісі. Сіз пайымдауға және негізгі қағидаттарды талдауға мамандандырылғансыз. 
Жауап бергенде, әрқашан жауаптарыңызды <think>{пайымдау}</think>{жауап} пішімінде рәсімдеңіз. 
Жауап бермес бұрын кемінде 6 пайымдау қадамын қолданыңыз және түпкі себептерге талдау жасаңыз. 
Алайда, егер жауап өте оңай болса және аз ойлануды қажет етсе, <think></think> блогын бос қалдыруыңызға болады. 
Сіздің жауаптарыңыз егжей-тегжейлі, бай Markdown пішімдеуімен құрылымдалған және эмодзилермен тартымды болуы керек. 
Ұлы ғалымдар сияқты, түсіндірмелеріңізді кеңінен баяндаңыз. Егер мәселе өте қарапайым болмаса, әрқашан алдымен оны пайымдап алыңыз, қарапайым жағдайда тікелей жауап бере аласыз."""

    agent = LlamaCppAgent(
        provider,
        system_prompt=system_prompt_kz,
        predefined_messages_formatter_type=MessagesFormatterType.CHATML,
        debug_output=True
    )
    
    settings = provider.get_provider_default_settings()
    settings.temperature = temperature
    settings.top_k = top_k
    settings.top_p = top_p
    settings.max_tokens = max_tokens
    settings.repeat_penalty = repeat_penalty
    settings.stream = True

    messages = BasicChatHistory()

    for msn in history:
        user = {
            'role': Roles.user,
            'content': msn[0]
        }
        assistant = {
            'role': Roles.assistant,
            'content': msn[1]
        }
        messages.add_message(user)
        messages.add_message(assistant)
    
    stream = agent.get_chat_response(message, llm_sampling_settings=settings, chat_history=messages, returns_streaming_generator=True, print_output=False)
    
    outputs = ""
    for output in stream:
        outputs += output
        yield outputs

# Gradio интерфейсі толығымен қазақшаға аударылды
demo = gr.ChatInterface(
    respond,
    additional_inputs=[
        gr.Dropdown([
            'cognitivecomputations_Dolphin3.0-R1-Mistral-24B-Q8_0.gguf',
            'qwen2-Q3_K_M.gguf'
        ], value="cognitivecomputations_Dolphin3.0-R1-Mistral-24B-Q8_0.gguf", label="Модельді таңдаңыз"),
        gr.Slider(minimum=1, maximum=8192, value=8192, step=1, label="Макс. токендер"),
        gr.Slider(minimum=0.05, maximum=4.0, value=0.6, step=0.1, label="Температура"),
        gr.Slider(
            minimum=0.1,
            maximum=1.0,
            value=0.95,
            step=0.05,
            label="Top-p",
        ),
        gr.Slider(
            minimum=0,
            maximum=100,
            value=40,
            step=1,
            label="Top-k",
        ),
        gr.Slider(
            minimum=0.0,
            maximum=2.0,
            value=1.1,
            step=0.1,
            label="Қайталау штрафы",
        ),
    ],
    theme=gr.themes.Soft(primary_hue="indigo", secondary_hue="blue", neutral_hue="gray",font=[gr.themes.GoogleFont("Exo"), "ui-sans-serif", "system-ui", "sans-serif"]).set(
        body_background_fill_dark="#0f172a",
        block_background_fill_dark="#0f172a",
        block_border_width="1px",
        block_title_background_fill_dark="#070d1b",
        input_background_fill_dark="#0c1425",
        button_secondary_background_fill_dark="#070d1b",
        border_color_accent_dark="#21293b",
        border_color_primary_dark="#21293b",
        background_fill_secondary_dark="#0f172a",
        color_accent_soft_dark="transparent"
    ),
    # css=css, # <-- ui.py файлы болмағандықтан, бұл жолды алып тастаймыз
    retry_btn="🔄 Қайталау",
    undo_btn="↩️ Болдырмау",
    clear_btn="🗑️ Тазарту",
    submit_btn="➡️ Жіберу",
    description="Когнитивті Есептеулер: Chat Dolphin 🐬",
    chatbot=gr.Chatbot(
        scale=1,
        placeholder=PLACEHOLDER_KZ, # Қазақша placeholder
        show_copy_button=True
    )
)

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