# app.py
import os
import re
import gradio as gr
import torch
from transformers import pipeline, AutoTokenizer

# ============================
# 3 Transformers (pipelines)
# ============================
# 1) Intent / zero-shot
DEFAULT_INTENT_MODEL = os.getenv("INTENT_MODEL", "joeddav/xlm-roberta-large-xnli")

# 2) Checklist generator (instruct)
DEFAULT_GEN_MODEL = os.getenv("GEN_MODEL", "Qwen/Qwen2.5-1.5B-Instruct")
# If your Space can handle it, this is often better structured:
# DEFAULT_GEN_MODEL = os.getenv("GEN_MODEL", "Qwen/Qwen2.5-1.5B-Instruct")

# 3) QA over checklist
DEFAULT_QA_MODEL=os.getenv("QA_MODEL", "MilyaShams/rubert-russian-qa-sberquad")

DEVICE = 0 if torch.cuda.is_available() else -1


def safe_make_pipeline(task: str, model_name: str, **kwargs):
    try:
        return pipeline(task, model=model_name, device=DEVICE, **kwargs), model_name
    except Exception:
        if task == "zero-shot-classification":
            fallback = "facebook/bart-large-mnli"
        elif task == "text-generation":
            fallback = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
        elif task == "question-answering":
            fallback = "distilbert-base-cased-distilled-squad"
        else:
            raise
        return pipeline(task, model=fallback, device=DEVICE, **kwargs), fallback


intent_pipe, intent_model_used = safe_make_pipeline(
    "zero-shot-classification",
    DEFAULT_INTENT_MODEL,
)
gen_pipe, gen_model_used = safe_make_pipeline(
    "text-generation",
    DEFAULT_GEN_MODEL,
)
qa_pipe, qa_model_used = safe_make_pipeline(
    "question-answering",
    DEFAULT_QA_MODEL,
)

# For Qwen chat formatting (works for many instruct LMs too)
try:
    gen_tokenizer = AutoTokenizer.from_pretrained(gen_model_used, use_fast=True)
except Exception:
    gen_tokenizer = None

DEFAULT_LABELS = [
    "обучение",
    "переезд",
    "релокация/иммиграция",
    "путешествие",
    "карьера/поиск работы",
    "финансы/покупка",
    "здоровье/фитнес",
    "ремонт/быт",
]

CATEGORY_CHOICES = ["Авто (определить по тексту)"] + DEFAULT_LABELS


def normalize_text(s: str) -> str:
    s = (s or "").strip()
    s = re.sub(r"\s+", " ", s)
    return s


def infer_intent(user_goal: str, labels: list[str]):
    if not user_goal:
        return "не задано", 0.0, "Нет входного текста."

    result = intent_pipe(user_goal, candidate_labels=labels, multi_label=False)
    top_label = result["labels"][0]
    top_score = float(result["scores"][0])

    lines = ["Распознавание намерения (zero-shot):"]
    for lab, sc in zip(result["labels"], result["scores"]):
        lines.append(f"- {lab}: {sc:.3f}")
    return top_label, top_score, "\n".join(lines)


def build_checklist_prompt(user_goal: str, theme: str | None, style: str, constraints: str) -> str:
    theme_part = f"Тема: {theme}\n" if theme else ""
    constraints_part = f"Контекст: {constraints}\n" if constraints else ""

    style_hint = {
        "кратко": "короткие пункты без лишних слов",
        "подробно": "чуть более подробные пункты + подпункты где уместно",
        "с акцентом на риски": "сильный акцент на предотвращение ошибок и рисков",
        "с акцентом на сроки": "добавляй дедлайны/временные окна там, где уместно",
    }.get(style, "короткие пункты")

    return (
        "Составь практичный чек-лист на русском языке.\n"
        "Верни ТОЛЬКО чек-лист без вступлений.\n"
        "Запрещено:\n"
        "- любые вступления/комментарии (например: 'Конечно!', 'Вот исправленный текст', 'не меняю смысл')\n"
        "- плейсхолдеры и шаблоны: никаких '[секунды]', '[какой-то текст]', '{...}', '<...>'\n"
        "- таблицы, 'Расчёт', 'Банк', 'Сообщение', поля для заполнения\n\n"
        "- любые сокращения и метки кроме (P0), (P1), (P2)\n"
        "- нельзя писать '(пос.)', '(пи)', '(n1)' и любые другие скобочные пометки\n"
        "Формат строго:\n"
        "- [ ] (P0) пункт\n"
        "  - подпункт (если нужно)\n"
        "Где P0 = срочно/критично, P1 = важно, P2 = можно позже.\n\n"
        "Требования:\n"
        f"- стиль: {style_hint}\n"
        "- 12–18 пунктов\n"
        "- без нумерации (никаких '1.'), только '- [ ]'\n"
        "- КАЖДЫЙ пункт должен начинаться с (P0) или (P1) или (P2)\n"
        "- Сначала выведи все P0, потом P1, потом P2\n"
        "- пункты конкретные и выполнимые\n"
        "- для пунктов про документы/банки/счета/визы/страховку ОБЯЗАТЕЛЬНО добавь 2–5 подпунктов\n"
        "- подпункты начинай с '  - '\n"
        "- если тема релокации/иммиграции: обязательно охвати документы, финансы, связь, жильё, медицину, язык\n"
        "- в конце 2 блока:\n"
        "- после приоритета НЕ добавляй другие метки в скобках\n"
        "- пример: '- [ ] (P1) Открыть банковский счёт'\n"
        "Проверка готовности:\n"
        "- ... (3–5 вопросов)\n"
        "Риски и как снизить:\n"
        "- ... (3–6 пунктов)\n\n"
        f"{theme_part}"
        f"{constraints_part}"
        f"Цель: {user_goal}\n\n"
        "Чек-лист:\n"
    )


def apply_chat_template_if_available(user_prompt: str) -> str:
    if gen_tokenizer is None or not hasattr(gen_tokenizer, "apply_chat_template"):
        return user_prompt
    messages = [
        {"role": "system", "content": "Ты пишешь грамотно по-русски и строго соблюдаешь формат чек-листа."},
        {"role": "user", "content": user_prompt},
    ]
    return gen_tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)


def clean_checklist(text: str) -> str:
    text = (text or "").strip()
    text = re.sub(r"<extra_id_\d+>", "", text).strip()

    cleaned = []
    current_group = "OTHER"

    for ln in text.splitlines():
        s = ln.rstrip()
        if not s.strip():
            continue

        low = s.strip().lower()

        bad_prefixes = (
            "конечно",
            "вот исправленный",
            "вот исправленный текст",
            "исправленный текст",
            "ничего не добавляя",
            "не меняю смысл",
            "я не меняю смысл",
            "готово",
            "итог",
            "результат",
        )
        if low.startswith(bad_prefixes):
            continue

        if low.startswith(("расч", "банк", "сообщение", "время ", "план поездки", "расход")):
            continue

        s2 = re.sub(r"\[(?!\s*\])([^\]]+)\]", "", s).strip()

        s2 = re.sub(r"^\-\s*\[\s*\]\s*\[(.+)\]\s*$", r"- [ ] \1", s2).strip()

        s2 = s2.replace("•", "-").replace("–", "-")
        s2 = re.sub(r"\s+", " ", s2)

        if s2.startswith("- [ ]") or s2.startswith("  -") or s2.startswith("Проверка готовности") or s2.startswith("Риски"):
            cleaned.append(s2)

    out = "\n".join(cleaned).strip()

    if "- [ ]" not in out:
        raw = [l.strip() for l in text.splitlines() if l.strip()]
        raw = [re.sub(r"^\d+[\).\s]+", "", r).strip(" -•\t") for r in raw]
        raw = [r for r in raw if r]
        raw = [f"- [ ] (P1) {r}" for r in raw[:18]]
        out = "\n".join(raw).strip()

    return out


def split_tail_sections(text: str):
    lines = text.splitlines()
    idx = None
    for i, ln in enumerate(lines):
        if ln.strip().startswith("Проверка готовности"):
            idx = i
            break
    if idx is None:
        return text, ""
    return "\n".join(lines[:idx]).strip(), "\n".join(lines[idx:]).strip()


def sort_by_priority(text: str) -> str:
    body, tail = split_tail_sections(text)
    if not body:
        return text

    groups = {"P0": [], "P1": [], "P2": [], "OTHER": []}
    current = "OTHER"

    for ln in body.splitlines():
        s = ln.strip()
        if s.startswith("- [ ]"):
            if "(P0)" in s:
                current = "P0"
            elif "(P1)" in s:
                current = "P1"
            elif "(P2)" in s:
                current = "P2"
            else:
                current = "OTHER"
            groups[current].append(ln)
        else:
            groups[current].append(ln)

    out_lines = []
    for k in ["P0", "P1", "P2", "OTHER"]:
        if groups[k]:
            out_lines.extend(groups[k])

    out = "\n".join(out_lines).strip()
    if tail:
        out = (out + "\n\n" + tail).strip()
    return out

def dedupe_checklist(text: str) -> str:
    body, tail = split_tail_sections(text)

    lines = body.splitlines()
    blocks = []
    i = 0
    while i < len(lines):
        if lines[i].startswith("- [ ]"):
            item = lines[i]
            subs = []
            i += 1
            while i < len(lines) and lines[i].startswith("  -"):
                subs.append(lines[i])
                i += 1
            blocks.append((item, subs))
        else:
            i += 1

    def norm(s: str) -> str:
        s = s.lower()
        s = re.sub(r"\(p[0-2]\)", "", s)
        s = re.sub(r"[^a-zа-я0-9 ]+", " ", s)
        s = re.sub(r"\s+", " ", s).strip()
        return s

    seen_items = set()
    out = []
    for item, subs in blocks:
        k = norm(item)
        if k in seen_items:
            continue
        seen_items.add(k)

        out.append(item)

        seen_subs = set()
        for sub in subs:
            sk = norm(sub)
            if sk in seen_subs:
                continue
            seen_subs.add(sk)
            out.append(sub)

    out_text = "\n".join(out).strip()
    if tail:
        out_text = (out_text + "\n\n" + tail).strip()
    return out_text

def trim_incomplete_last_line(text: str) -> str:
    lines = text.splitlines()
    if not lines:
        return text

    last = lines[-1].strip()

    if (last.startswith("- [ ]") or last.startswith("  -")) and (
        re.search(r"[,:;—–-]$", last) or len(last.split()) < 3
    ):
        lines = lines[:-1]

    return "\n".join(lines).strip()


def ensure_doc_finance_subpoints(text: str) -> str:
    lines = text.splitlines()
    out = []

    added_docs = False
    added_fin = False

    def has_any_subpoints_with_keywords(all_lines, keywords):
        low_all = "\n".join(all_lines).lower()
        return any(k in low_all for k in keywords)

    docs_kw_present = has_any_subpoints_with_keywords(lines, ["загранпаспорт", "виза", "внж", "копии", "облако", "приглашение"])
    fin_kw_present = has_any_subpoints_with_keywords(lines, ["комисси", "лимит", "вторая карта", "2fa", "двухфактор", "уведомлен"])

    i = 0
    while i < len(lines):
        ln = lines[i]
        out.append(ln)

        if ln.startswith("- [ ]"):
            low = ln.lower()

            j = i + 1
            has_sub = (j < len(lines) and lines[j].startswith("  -"))

            if (not added_docs) and (not docs_kw_present) and (not has_sub) and any(k in low for k in ["документ", "виза", "внж", "страхов", "паспорт"]):
                out.extend([
                    "  - Проверь срок действия загранпаспорта и требования к остаточному сроку",
                    "  - Составь список нужных документов и сделай копии (бумага + облако)",
                    "  - Уточни требования по визе/ВНЖ и собери подтверждения (финансы, жильё, приглашение)",
                ])
                added_docs = True

            if (not added_fin) and (not fin_kw_present) and (not has_sub) and any(k in low for k in ["счёт", "счет", "банк", "карта", "финанс"]):
                out.extend([
                    "  - Проверь лимиты, комиссии и возможность работы карт за границей",
                    "  - Подготовь резервный доступ к деньгам (вторая карта/наличные/перевод)",
                    "  - Включи уведомления и двухфакторную защиту в банковских приложениях",
                ])
                added_fin = True

        i += 1

    return "\n".join(out).strip()


def polish_russian_same_model(checklist_text: str) -> str:
    if not checklist_text or checklist_text.count("- [ ]") < 6:
        return checklist_text

    polish_prompt = (
        "Отредактируй текст чек-листа: исправь грамматику и стиль русского языка.\n"
        "ВАЖНО:\n"
        "- Верни ТОЛЬКО чек-лист.\n"
        "- Никаких вступлений (например: 'Конечно', 'Вот исправленный текст').\n"
        "- Никаких комментариев (например: 'не меняю смысл').\n"
        "- Ничего не добавляй и не удаляй по смыслу.\n"
        "- Сохрани формат:\n"
        "  - [ ] (P0/P1/P2) пункт\n"
        "    - подпункт\n"
        "- Сохрани блоки 'Проверка готовности' и 'Риски и как снизить'.\n\n"
        "- не меняй формулировки на канцелярит, пиши простым естественным русским\n"
        f"{checklist_text}\n"
    )
    polish_prompt = apply_chat_template_if_available(polish_prompt)

    outp = gen_pipe(
        polish_prompt,
        max_new_tokens=260,
        do_sample=False,
        temperature=0.0,  # если поддерживается
        return_full_text=False,
    )
    text2 = (outp[0].get("generated_text") or "").strip()
    text2 = clean_checklist(text2)
    text2 = sort_by_priority(text2)

    if text2.count("- [ ]") >= checklist_text.count("- [ ]"):
        return text2
    return checklist_text


def generate_checklist(user_goal: str, category: str, style: str, constraints: str):
    user_goal = normalize_text(user_goal)
    constraints = normalize_text(constraints)

    if not user_goal:
        return (
            "",
            "Введите описание цели.",
            "",
            f"intent={intent_model_used}\ngen={gen_model_used}\nqa={qa_model_used}",
            None,
            None,
        )

    inferred_label, inferred_score, intent_debug = infer_intent(user_goal, DEFAULT_LABELS)

    if category and category != "Авто (определить по тексту)":
        chosen_theme = category
    else:
        chosen_theme = inferred_label if inferred_score >= 0.30 else None

    base_prompt = build_checklist_prompt(
        user_goal=user_goal,
        theme=chosen_theme,
        style=style,
        constraints=constraints,
    )
    prompt = apply_chat_template_if_available(base_prompt)

    out = gen_pipe(
        prompt,
        max_new_tokens=900,
        do_sample=True,
        temperature=0.55,
        top_p=0.9,
        repetition_penalty=1.15,
        no_repeat_ngram_size=4,
        return_full_text=False,
    )
    text = (out[0].get("generated_text") or "").strip()
    text = clean_checklist(text)
    text = sort_by_priority(text)
    text = ensure_doc_finance_subpoints(text)
    text = dedupe_checklist(text)
    text = trim_incomplete_last_line(text)
    text = polish_russian_same_model(text)
    text = dedupe_checklist(text)
    text = trim_incomplete_last_line(text)


    if text.count("- [ ]") < 10:
        retry_prompt = (
            "Сделай чек-лист строго в формате '- [ ] (P0/P1/P2) ...' (12–18 пунктов) на русском.\n"
            "Сначала P0, затем P1, затем P2.\n"
            "Без вступлений. Без комментариев. Без плейсхолдеров в квадратных скобках.\n"
            "Для документов/банков/счётов/визы/страховки добавь подпункты (2–5).\n"
            "В конце добавь:\n"
            "Проверка готовности: (3–5 вопросов)\n"
            "Риски и как снизить: (3–6 пунктов)\n\n"
            f"Цель: {user_goal}\n"
            f"Контекст: {constraints}\n"
            f"Тема: {chosen_theme}\n\n"
            "Чек-лист:\n"
        )
        retry_prompt = apply_chat_template_if_available(retry_prompt)
        out2 = gen_pipe(
            retry_prompt,
            max_new_tokens=750,
            do_sample=True,
            temperature=0.75,
            top_p=0.93,
            repetition_penalty=1.12,
            return_full_text=False,
        )
        text2 = (out2[0].get("generated_text") or "").strip()
        text2 = clean_checklist(text2)
        text2 = sort_by_priority(text2)
        if text2.count("- [ ]") > text.count("- [ ]"):
            text = text2

    text = ensure_doc_finance_subpoints(text)

    text = polish_russian_same_model(text)

    if text.count("- [ ]") < 8:
        text = (
            "- [ ] (P0) Не удалось корректно сформировать чек-лист.\n"
            "- [ ] (P0) Попробуйте уточнить цель (страна/город, статус визы, бюджет, сроки) и нажать ещё раз.\n"
        )

    meta = {
        "goal": user_goal,
        "theme": chosen_theme,
        "intent_label": inferred_label,
        "intent_score": inferred_score,
        "intent_model": intent_model_used,
        "gen_model": gen_model_used,
        "qa_model": qa_model_used,
    }

    models_info = f"intent={intent_model_used}\ngen={gen_model_used}\nqa={qa_model_used}"
    theme_info = f"{chosen_theme or '(не определилась)'}"

    return text, intent_debug, theme_info, models_info, text, meta


def answer_question(question: str, checklist_state: str, meta_state: dict | None):
    question = normalize_text(question)
    if not checklist_state:
        return "Сначала сгенерируйте чек-лист на первой вкладке.", ""

    if not question:
        return "Введите вопрос.", ""

    qa_res = qa_pipe(question=question, context=checklist_state)
    answer = (qa_res.get("answer") or "").strip()
    score = float(qa_res.get("score") or 0.0)

    evidence = f"QA score: {score:.3f}\n"
    if answer:
        evidence += f"Span: {answer}\n"

    if (not answer) or score < 0.20 or len(answer) < 3:
        goal = (meta_state or {}).get("goal", "")
        theme = (meta_state or {}).get("theme", "")

        user_prompt = (
            "Ответь на вопрос по чек-листу. Пиши по-русски, кратко и практично.\n"
            "Если ответа нет в чек-листе — предложи 3–6 дополнительных пунктов (с приоритетом P0/P1/P2).\n"
            "Не используй вступления и комментарии.\n"
            "Не используй плейсхолдеры в квадратных скобках.\n\n"
            f"Цель: {goal}\n"
            f"Тема: {theme}\n\n"
            f"Чек-лист:\n{checklist_state}\n\n"
            f"Вопрос: {question}\n"
            "Ответ:\n"
        )
        prompt = apply_chat_template_if_available(user_prompt)
        gen_out = gen_pipe(
            prompt,
            max_new_tokens=320,
            do_sample=False,
            return_full_text=False,
        )[0]["generated_text"].strip()
        gen_out = re.sub(r"<extra_id_\d+>", "", gen_out).strip()
        return gen_out, evidence + "Fallback: generator used (QA confidence low)."

    return f"{answer}\n\n_(Найдено в чек-листе; уверенность: {score:.2f})_", evidence



with gr.Blocks(title="Умный чек-лист (3 Transformers)") as demo:
    gr.Markdown(
        "# ✅ Умный чек-лист (3 Transformers)\n"
        "1) распознаём намерение → 2) генерируем чек-лист → 3) отвечаем на вопросы по чек-листу\n"
    )

    checklist_state = gr.State(value=None)
    meta_state = gr.State(value=None)

    with gr.Tab("1) Создать чек-лист"):
        with gr.Row():
            with gr.Column(scale=3):
                user_goal = gr.Textbox(
                    label="Опишите, что вы хотите сделать",
                    placeholder="Напр.: Хочу переехать в Китай на 2 года. Бюджет 200 000 ₽. Есть загранпаспорт.",
                    lines=4,
                )

                category = gr.Dropdown(
                    label="Категория (необязательно)",
                    choices=CATEGORY_CHOICES,
                    value="Авто (определить по тексту)",
                )

                style = gr.Dropdown(
                    label="Стиль чек-листа",
                    choices=["кратко", "подробно", "с акцентом на риски", "с акцентом на сроки"],
                    value="кратко",
                )

                constraints = gr.Textbox(
                    label="Контекст/ограничения (необязательно)",
                    placeholder="Напр.: бюджет, срок, город/страна, семья/дети, удалённая работа, уровень языка...",
                    lines=3,
                )

                gen_btn = gr.Button("Сгенерировать чек-лист", variant="primary")

            with gr.Column(scale=2):
                theme_box = gr.Textbox(
                    label="Выбранная/распознанная тема",
                    lines=2,
                    interactive=False,
                )
                models_box = gr.Textbox(
                    label="Модели",
                    lines=3,
                    interactive=False,
                )
                intent_debug = gr.Textbox(
                    label="Диагностика намерения",
                    lines=10,
                    interactive=False,
                )

        checklist_out = gr.Code(label="Чек-лист", language="markdown")

        gen_btn.click(
            fn=generate_checklist,
            inputs=[user_goal, category, style, constraints],
            outputs=[checklist_out, intent_debug, theme_box, models_box, checklist_state, meta_state],
        )

    with gr.Tab("2) Уточняющие вопросы по чек-листу"):
        with gr.Row():
            with gr.Column(scale=3):
                gr.Markdown("Задайте вопрос по уже сгенерированному чек-листу.")
                question = gr.Textbox(
                    label="Ваш вопрос",
                    placeholder="Напр.: Какие документы подготовить? Как снизить расходы? Что сделать первым делом?",
                    lines=2,
                )
                ask_btn = gr.Button("Ответить", variant="primary")
                answer_out = gr.Markdown(label="Ответ")
            with gr.Column(scale=2):
                evidence_out = gr.Textbox(label="Тех. детали", lines=10)

        ask_btn.click(
            fn=answer_question,
            inputs=[question, checklist_state, meta_state],
            outputs=[answer_out, evidence_out],
        )

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