import os
import re
import html as ihtml
import textwrap
import warnings
import requests
import gradio as gr

from smolagents import CodeAgent, LiteLLMModel, tool

# Убираем шумные предупреждения (на работу не влияет)
warnings.filterwarnings("ignore", category=UserWarning, module="pydantic")


# ----------------------------
# Config (CPU-friendly defaults)
# ----------------------------
OLLAMA_BASE = os.getenv("OLLAMA_URL", "http://127.0.0.1:11434").rstrip("/")

# На CPU 30B почти всегда будет очень медленно/таймауты.
# Рекомендуемый дефолт: qwen2.5-coder:1.5b или 3b
MODEL_NAME = os.getenv("MODEL_NAME", "qwen2.5-coder:3b")

TRIAGE_MODEL = os.getenv("TRIAGE_MODEL", MODEL_NAME)
ACTIONS_MODEL = os.getenv("ACTIONS_MODEL", MODEL_NAME)
WRITER_MODEL = os.getenv("WRITER_MODEL", MODEL_NAME)

# На CPU держим разумные лимиты
NUM_CTX = int(os.getenv("NUM_CTX", "2048"))
MAX_TOKENS = int(os.getenv("MAX_TOKENS", "768"))
TIMEOUT = int(os.getenv("LITELLM_TIMEOUT", "3600"))


def make_model(model_name: str) -> LiteLLMModel:
    # flatten_messages_as_text=False помогает избежать проблем формата messages в некоторых связках smolagents+litellm+ollama
    return LiteLLMModel(
        model_id=f"ollama_chat/{model_name}",
        api_base=OLLAMA_BASE,
        num_ctx=NUM_CTX,
        temperature=0.2,
        max_tokens=MAX_TOKENS,
        timeout=TIMEOUT,
        flatten_messages_as_text=False,
    )


def _strip_html(raw_html: str) -> str:
    """Грубая очистка HTML -> текст."""
    raw_html = ihtml.unescape(raw_html)
    raw_html = re.sub(r"(?is)<(script|style).*?>.*?</\1>", " ", raw_html)
    raw_html = re.sub(r"(?is)<br\s*/?>", "\n", raw_html)
    raw_html = re.sub(r"(?is)</p\s*>", "\n", raw_html)
    raw_html = re.sub(r"(?is)<.*?>", " ", raw_html)
    raw_html = re.sub(r"[ \t]+", " ", raw_html)
    raw_html = re.sub(r"\n{3,}", "\n\n", raw_html)
    return raw_html.strip()


@tool
def web_search(query: str) -> str:
    """Поиск (DuckDuckGo HTML) -> очищенный текст.

    Args:
        query: Текстовый поисковый запрос.

    Returns:
        Очищенный текст (до ~3500 символов).
    """
    resp = requests.get(
        "https://duckduckgo.com/html/",
        params={"q": query},
        timeout=12,
        headers={"User-Agent": "financial-aid-navigator-demo"},
    )
    resp.raise_for_status()
    return _strip_html(resp.text)[:3500]


@tool
def extract_facts(case_description: str) -> str:
    """Извлекает факты из описания без домыслов.

    Args:
        case_description: Текст пользователя.

    Returns:
        Факты + что уточнить (безопасно).
    """
    text = (case_description or "").strip()
    low = text.lower()

    flags = []
    if re.search(r"\b(просроч|просрочка|пропустил платеж)\b", low):
        flags.append("просрочка/пропуск платежа")
    if re.search(r"\b(коллектор|взыскател)\b", low):
        flags.append("контакт/визит коллекторов")
    if re.search(r"\b(ипотек)\b", low):
        flags.append("ипотека")
    if re.search(r"\b(кредит|займ|мфо)\b", low):
        flags.append("кредит/займ/МФО")

    nums = re.findall(r"\d[\d\s]{2,}", text)
    nums = [re.sub(r"\s+", "", n) for n in nums][:8]

    missing = [
        "Сколько дней/месяцев просрочка?",
        "Какая сумма просрочки/долга? (можно диапазоном)",
        "Кто кредитор/банк/МФО? (без номеров договоров)",
        "Есть ли письма/суды/исполнительное производство?",
        "Есть ли угрозы/давление (звонки родственникам, визиты)?",
        "Какой официальный доход/обязательные расходы? (диапазоны)",
    ]

    return (
        "ФАКТЫ ИЗ ОПИСАНИЯ (без домыслов):\n"
        f"- Исходный текст: {text}\n"
        f"- Признаки: {', '.join(flags) if flags else 'не выявлено'}\n"
        f"- Числа: {', '.join(nums) if nums else 'нет'}\n"
        f"- Что уточнить (безопасно): {'; '.join(missing)}\n"
        "Правило: нельзя добавлять факты, которых нет в исходном тексте."
    )


# В вашей версии CodeAgent парсит код из <code>...</code>, поэтому используем именно этот формат.
# Мы ставим max_steps=1 (быстрее), поэтому сильно ужесточаем формат — чтобы модель не “съехала”.
CODE_FORMAT_RULES = """
КРИТИЧЕСКИ ВАЖНО: это CodeAgent.
Ты обязан вернуть ТОЛЬКО Python-код внутри тегов <code>...</code>.
Ответ должен НАЧИНАТЬСЯ с "<code>" и ЗАКАНЧИВАТЬСЯ "</code>".
Никакого текста до или после <code>...</code>.
В конце кода обязательно вызови final_answer(...) и больше ничего.
"""


def _friendly_agent_error(where: str, err: Exception) -> str:
    return (
        f"### Ошибка на шаге: {where}\n\n"
        "Модель не вернула корректный блок кода для CodeAgent или код не выполнился.\n\n"
        "Быстрые способы починить:\n"
        "- Поставить модель поменьше: `MODEL_NAME=qwen2.5-coder:1.5b`\n"
        "- Уменьшить `MAX_TOKENS` до 512–768\n"
        "- Оставить `NUM_CTX` 2048\n\n"
        "Текст ошибки:\n"
        f"```text\n{repr(err)}\n```"
    )


def run_fin_aid_multi_agent(case_description: str, region: str, urgency: str, allow_internet: bool):
    case_description = (case_description or "").strip()
    region = (region or "").strip()
    urgency = (urgency or "").strip()

    if not case_description:
        return "Пожалуйста, заполните поле «Описание ситуации»."

    # Для скорости и стабильности:
    # - всегда доступен extract_facts
    # - web_search доступен только если allow_internet=True
    tools = [extract_facts] + ([web_search] if allow_internet else [])

    # ============================
    # TRIAGE (max_steps=1)
    # ============================
    triage_agent = CodeAgent(
        tools=tools,
        model=make_model(TRIAGE_MODEL),
        add_base_tools=False,
        max_steps=1,  # БЫСТРО: один шаг, без повторных попыток
        additional_authorized_imports=[],
    )

    triage_prompt = f"""
Ты агент TRIAGE (финансовая помощь). Всегда по-русски.
{CODE_FORMAT_RULES}

Задача: быстро разобрать ситуацию и приоритеты.

Правила:
- НЕ выдумывай факты. Используй только входные данные и extract_facts().
- НЕ предполагай “взлом/мошенничество”, если пользователь этого не писал.
- НЕ проси номера карт/CVV/пароли/SMS-коды/номера документов.

Если web_search доступен (проверь: 'web_search' in globals()):
- сделай ОДИН (1) вызов web_search по запросу:
  web_search(f"emergency financial assistance {region}")
(Для скорости: только один запрос. Не вставляй длинные куски — максимум 3–5 тезисов.)

Код ДОЛЖЕН начинаться с объявления переменных (буквально):
case_description = \"\"\"...\"\"\"
region = \"\"\"...\"\"\"
urgency = \"\"\"...\"\"\"

Далее:
facts = extract_facts(case_description)
(опционально) s1 = web_search(...)
triage_text = \"\"\"...\"\"\"  # Markdown

Структура triage_text:
## Сводка ситуации
## Приоритеты
- Сегодня (24–72 часа)
- На неделе
- В течение месяца
## Риски
## Какие данные подготовить (безопасно)
## Вопросы для уточнения (до 8)

В конце: final_answer(triage_text)

Входные данные:
urgency = \"\"\"{urgency}\"\"\"
region = \"\"\"{region}\"\"\"
case_description = \"\"\"{case_description}\"\"\"
"""
    try:
        triage_result = triage_agent.run(textwrap.dedent(triage_prompt))
    except Exception as e:
        return _friendly_agent_error("TRIAGE", e)

    # ============================
    # ACTIONS (max_steps=1)
    # ============================
    actions_agent = CodeAgent(
        tools=tools,
        model=make_model(ACTIONS_MODEL),
        add_base_tools=False,
        max_steps=1,  # БЫСТРО
        additional_authorized_imports=[],
    )

    actions_prompt = f"""
Ты агент ACTIONS (финансовая помощь). Всегда по-русски.
{CODE_FORMAT_RULES}

Правила:
- НЕ выдумывай факты. Опирайся на triage_result и extract_facts(case_description).
- НЕ проси номера карт/CVV/пароли/SMS-коды/номера документов.
- Если кейс про просрочку/коллекторов: приоритет — деэскалация, фиксация общения, законные шаги, переговоры с кредитором.

Код ДОЛЖЕН начинаться с объявления переменных (буквально):
case_description = \"\"\"...\"\"\"
region = \"\"\"...\"\"\"
urgency = \"\"\"...\"\"\"

Далее:
facts = extract_facts(case_description)
actions_text = \"\"\"...\"\"\"  # Markdown

Структура actions_text:
## Действия
### Срочно (urgent)
(4–7 пунктов)
### На неделе (short_term)
(4–7 пунктов)
### В течение месяца (mid_term)
(3–5 пунктов)

## Варианты помощи в регионе
(6–10 категорий: соцзащита/муниципалитет/НКО/юристы по долгам/банк/кредитные каникулы и т.п.)

## Анти-мошенничество
(6–8 пунктов: “не платить предоплату”, “не давать коды/пароли”, “проверять юрлицо”, и т.д.)

## Уточняющие вопросы (до 8)

В конце: final_answer(actions_text)

Входные данные:
urgency = \"\"\"{urgency}\"\"\"
region = \"\"\"{region}\"\"\"
case_description = \"\"\"{case_description}\"\"\"

triage_result:
{triage_result}
"""
    try:
        actions_result = actions_agent.run(textwrap.dedent(actions_prompt))
    except Exception as e:
        return _friendly_agent_error("ACTIONS", e)

    # ============================
    # WRITER (FINAL) (max_steps=1)
    # ============================
    writer_agent = CodeAgent(
        tools=tools,
        model=make_model(WRITER_MODEL),
        add_base_tools=False,
        max_steps=1,  # БЫСТРО
        additional_authorized_imports=[],
    )

    writer_prompt = f"""
Ты агент WRITER (финансовая помощь). Всегда по-русски.
{CODE_FORMAT_RULES}

Правила:
- НЕ выдумывай факты.
- НЕ проси номера карт/CVV/пароли/SMS-коды/номера документов.
- Не делай таблицу на 30 строк — только краткий мини-бюджет.

Код ДОЛЖЕН начинаться с объявления переменных (буквально):
case_description = \"\"\"...\"\"\"
region = \"\"\"...\"\"\"
urgency = \"\"\"...\"\"\"

Далее:
facts = extract_facts(case_description)
report = \"\"\"...\"\"\"  # Markdown
final_answer(report)

Структура report:
# План финансовой помощи
## Важно
## Сводка ситуации
## Приоритеты
### Сегодня (24–72 часа)
### На неделе
### В течение месяца
## Пошаговый план
## Варианты помощи в регионе
## Мини-бюджет на 30 дней
## Анти-мошенничество
## Что подготовить
## Вопросы для уточнения

Входные данные:
urgency = \"\"\"{urgency}\"\"\"
region = \"\"\"{region}\"\"\"
case_description = \"\"\"{case_description}\"\"\"

triage_result:
{triage_result}

actions_result:
{actions_result}
"""
    try:
        final_report = writer_agent.run(textwrap.dedent(writer_prompt))
        return str(final_report).strip() if final_report is not None else ""
    except Exception as e:
        return _friendly_agent_error("WRITER", e)


with gr.Blocks(title="Financial Aid Navigator (CodeAgent + Ollama)") as demo:
    gr.Markdown("# Financial Aid Navigator — Multi-agent (CodeAgent + Ollama)")
    gr.Markdown(
        f"- Ollama: `{OLLAMA_BASE}`\n"
        f"- MODEL_NAME: `{MODEL_NAME}`\n"
        f"- TRIAGE_MODEL: `{TRIAGE_MODEL}`\n"
        f"- ACTIONS_MODEL: `{ACTIONS_MODEL}`\n"
        f"- WRITER_MODEL: `{WRITER_MODEL}`\n"
        f"- NUM_CTX: `{NUM_CTX}`, MAX_TOKENS: `{MAX_TOKENS}`, TIMEOUT: `{TIMEOUT}`\n"
        f"- max_steps у агентов: `1` (быстрее)"
    )

    region = gr.Textbox(label="Регион", lines=1, placeholder="Например: Россия")
    urgency = gr.Dropdown(
        ["срочно (24–72 часа)", "в течение недели", "не срочно (в течение месяца)"],
        value="в течение недели",
        label="Срочность",
    )
    case_description = gr.Textbox(
        label="Описание ситуации",
        lines=8,
        placeholder="Например: Просрочил платеж по ипотеке, коллекторы стучатся в дверь. Сумма ...",
    )
    allow_internet = gr.Checkbox(label="Разрешить web_search (интернет)", value=False)

    run_btn = gr.Button("Сформировать план помощи (мультиагент)")
    output = gr.Textbox(label="Результат (Markdown)", lines=26)

    run_btn.click(
        fn=run_fin_aid_multi_agent,
        inputs=[case_description, region, urgency, allow_internet],
        outputs=[output],
        queue=False,
    )


def main():
    demo.launch(
        server_name=os.getenv("GRADIO_SERVER_NAME", "0.0.0.0"),
        server_port=int(os.getenv("GRADIO_SERVER_PORT", "7860")),
        show_error=True,
    )


if __name__ == "__main__":
    main()