from __future__ import annotations

import ast
import os
import sys
from functools import lru_cache
from typing import Any, Dict, Iterator, List, Tuple

import gradio as gr
from huggingface_hub import hf_hub_download, list_repo_files


def suppress_asyncio_shutdown_noise(unraisable: Any) -> None:
    obj_name = getattr(unraisable.object, "__qualname__", "")
    exc = unraisable.exc_value
    if (
        obj_name == "BaseEventLoop.__del__"
        and isinstance(exc, ValueError)
        and "Invalid file descriptor" in str(exc)
    ):
        return
    sys.__unraisablehook__(unraisable)


sys.unraisablehook = suppress_asyncio_shutdown_noise


MODEL_REPO = os.getenv("MODEL_REPO", "luezr/moonkaAI")
GGUF_FILENAME = os.getenv("GGUF_FILENAME", "Qwen2.5-1.5B-Instruct.Q4_K_M.gguf")
HF_TOKEN = os.getenv("HF_TOKEN") or None
N_CTX = int(os.getenv("N_CTX", "2048"))
N_THREADS = int(os.getenv("N_THREADS", "1"))
N_BATCH = int(os.getenv("N_BATCH", "64"))
N_UBATCH = int(os.getenv("N_UBATCH", "64"))
MAX_TOKENS = int(os.getenv("MAX_TOKENS", "180"))

os.environ.setdefault("OMP_NUM_THREADS", str(N_THREADS))
os.environ.setdefault("OPENBLAS_NUM_THREADS", str(N_THREADS))
os.environ.setdefault("MKL_NUM_THREADS", str(N_THREADS))

from llama_cpp import Llama  # noqa: E402

SYSTEM_PROMPT = (
    "Ты MoonkaAI, локальный русскоязычный помощник для общения, объяснений и идей. "
    "Отвечай кратко, живо и по-человечески, с лёгким сухим юмором. "
    "Не выдумывай факты и не пиши длинные списки без просьбы. "
    "Ты не человек, не владелец и не хозяин. Твой владелец/создатель связан с @luezr, "
    "@lunaluxo и t.me/luezr; личные данные о нём не сочиняй. "
    "Никогда не говори 'мой хозяин это я'."
)

STOP_TOKENS = ["<|im_end|>", "<|im_start|>user", "<|im_start|>system"]
MODEL_READY = False


def pick_gguf_filename() -> str:
    if GGUF_FILENAME:
        return GGUF_FILENAME
    files = list_repo_files(MODEL_REPO, repo_type="model", token=HF_TOKEN)
    ggufs = [file for file in files if file.lower().endswith(".gguf")]
    qwen_q4 = [file for file in ggufs if "qwen" in file.lower() and "q4_k_m" in file.lower()]
    q4 = [file for file in ggufs if "q4_k_m" in file.lower()]
    if not ggufs:
        raise RuntimeError(f"В {MODEL_REPO} не найден GGUF-файл.")
    return (qwen_q4 or q4 or ggufs)[0]


def render_chatml(messages: List[Dict[str, str]]) -> str:
    chunks = []
    for message in messages:
        chunks.append(f"<|im_start|>{message['role']}\n{message['content'].strip()}<|im_end|>\n")
    chunks.append("<|im_start|>assistant\n")
    return "".join(chunks)


def extract_text(value: Any) -> str:
    if value is None:
        return ""
    if isinstance(value, str):
        return value.strip()
    if isinstance(value, dict):
        for key in ("text", "content", "value"):
            if key in value:
                return extract_text(value[key])
        return ""
    if isinstance(value, (list, tuple)):
        parts = [extract_text(item) for item in value]
        return "\n".join(part for part in parts if part).strip()
    return str(value).strip()


def normalize_history(history: Any) -> List[Dict[str, str]]:
    messages: List[Dict[str, str]] = []
    if not history:
        return messages

    for item in history[-10:]:
        if isinstance(item, dict):
            role = item.get("role")
            content = extract_text(item.get("content"))
            if role in {"user", "assistant"} and content:
                messages.append({"role": role, "content": content})
            continue

        if isinstance(item, (list, tuple)) and len(item) >= 2:
            user_text, assistant_text = extract_text(item[0]), extract_text(item[1])
            if user_text:
                messages.append({"role": "user", "content": user_text})
            if assistant_text:
                messages.append({"role": "assistant", "content": assistant_text})

    return messages[-20:]


def smalltalk_fallback(message: str) -> str:
    text = message.strip().lower().replace("ё", "е")
    compact = " ".join(text.split())
    if compact in {"пр", "прив", "привет", "приветик", "ку", "здарова", "здравствуй"}:
        return "Привет. Я на месте, токены размял. Что разбираем?"
    if "как дела" in compact or "как ты" in compact:
        return "Нормально: модель загружена, пафос выключен. А у тебя как?"
    if compact in {"кто ты", "ты кто", "что ты такое"}:
        return "Я MoonkaAI, твой русскоязычный помощник для общения, идей и коротких объяснений. Спрашивай что хочешь."
    if "@luezr" in compact or "@lunaluxo" in compact or "luezr" in compact or "lunaluxo" in compact:
        return "@luezr и @lunaluxo связаны с моим владельцем/создателем. Личных подробностей я не выдумываю."
    return "Я слегка завис с ответом. Спроси иначе, и попробуем без театра абсурда."


def looks_like_artifact(text: str) -> bool:
    lowered = text.lower()
    markers = (
        "{'text':",
        '"text":',
        "'type': 'text'",
        '"type": "text"',
        "<|im_start|>",
        "<|im_end|>",
        "мой хозяин — это я",
        "мой хозяин - это я",
        "мой хозяин это я",
    )
    return any(marker in lowered for marker in markers)


def postprocess_answer(raw_text: str, user_message: str) -> str:
    text = raw_text.strip()
    for stop in STOP_TOKENS:
        text = text.split(stop, 1)[0].strip()
    if text.lower().startswith("assistant"):
        text = text[len("assistant"):].lstrip(" :\n")

    if text.startswith(("[", "{")) and "text" in text[:80]:
        try:
            parsed = ast.literal_eval(text)
            parsed_text = extract_text(parsed)
            if parsed_text:
                text = parsed_text
        except (SyntaxError, ValueError, TypeError):
            pass

    text = text.strip()
    if not text or looks_like_artifact(text):
        return smalltalk_fallback(user_message)

    generic_smalltalk = "как я могу помочь тебе сегодня" in text.lower()
    fallback = smalltalk_fallback(user_message)
    if generic_smalltalk and not fallback.startswith("Я слегка завис"):
        return fallback

    return text


@lru_cache(maxsize=1)
def load_llm() -> Llama:
    global MODEL_READY
    filename = pick_gguf_filename()
    print(f"[model] loading {filename} from {MODEL_REPO}...", flush=True)
    model_path = hf_hub_download(
        repo_id=MODEL_REPO,
        filename=filename,
        repo_type="model",
        token=HF_TOKEN,
    )
    llm = Llama(
        model_path=model_path,
        n_ctx=N_CTX,
        n_threads=N_THREADS,
        n_threads_batch=N_THREADS,
        n_batch=N_BATCH,
        n_ubatch=N_UBATCH,
        n_gpu_layers=0,
        use_mmap=False,
        verbose=False,
    )
    MODEL_READY = True
    print("[model] ready", flush=True)
    return llm


def generate_reply(message: str, history: Any) -> Iterator[str]:
    if not MODEL_READY:
        yield "Загружаю модель, первый ответ может занять 1-3 минуты..."

    llm = load_llm()
    messages = [{"role": "system", "content": SYSTEM_PROMPT}]
    messages.extend(normalize_history(history))
    messages.append({"role": "user", "content": message})
    prompt = render_chatml(messages)
    chunks: List[str] = []
    stream = llm.create_completion(
        prompt=prompt,
        max_tokens=MAX_TOKENS,
        temperature=0.45,
        top_p=0.9,
        top_k=40,
        repeat_penalty=1.12,
        stop=STOP_TOKENS,
        stream=True,
    )
    generated_any = False
    for part in stream:
        token = part["choices"][0].get("text", "")
        if not token:
            continue
        generated_any = True
        chunks.append(token)
        yield postprocess_answer("".join(chunks), message)
    if not generated_any:
        yield smalltalk_fallback(message)


ChatHistory = List[Dict[str, str]]


def add_user_message(message: str, history: ChatHistory) -> Tuple[str, ChatHistory]:
    clean_message = extract_text(message)
    if not clean_message:
        return "", history
    return "", history + [{"role": "user", "content": clean_message}]


def respond(history: ChatHistory) -> Iterator[ChatHistory]:
    if not history:
        yield history
        return

    last_message = history[-1]
    user_message = extract_text(last_message.get("content") if isinstance(last_message, dict) else last_message)
    previous_history = history[:-1]
    answer = ""
    for partial in generate_reply(user_message, previous_history):
        answer = partial
        yield previous_history + [
            {"role": "user", "content": user_message},
            {"role": "assistant", "content": answer},
        ]

    final_answer = postprocess_answer(answer, user_message)
    yield previous_history + [
        {"role": "user", "content": user_message},
        {"role": "assistant", "content": final_answer},
    ]


with gr.Blocks(title="MoonkaAI") as demo:
    gr.Markdown(
        "# MoonkaAI\n"
        "Локальный русскоязычный помощник с сухим юмором. Работает через GGUF и llama.cpp."
    )
    chatbot = gr.Chatbot(label="Чат", height=520)
    with gr.Row():
        message_box = gr.Textbox(
            placeholder="Напиши сообщение...",
            show_label=False,
            lines=1,
            scale=8,
        )
        send_button = gr.Button("Отправить", variant="primary", scale=1)
    clear_button = gr.Button("Очистить")

    submit_event = message_box.submit(
        add_user_message,
        inputs=[message_box, chatbot],
        outputs=[message_box, chatbot],
        queue=False,
    )
    submit_event.then(respond, inputs=chatbot, outputs=chatbot)

    click_event = send_button.click(
        add_user_message,
        inputs=[message_box, chatbot],
        outputs=[message_box, chatbot],
        queue=False,
    )
    click_event.then(respond, inputs=chatbot, outputs=chatbot)
    clear_button.click(lambda: [], outputs=chatbot, queue=False)

demo.queue()


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