app.py

# app.py
import os
import tempfile
from decimal import Decimal
from typing import List, Optional

import pandas as pd
from fastapi import FastAPI, HTTPException, UploadFile, File
from pydantic import BaseModel

from common import prepare_components_series, fit_and_forecast, current_month_snapshot
from advice import _load as advice_load, _gen as advice_gen, _to_bullets, clean_ru
from receipt_total_api import extract_total

app = FastAPI()

# ---------- Pydantic-модели под Go-структуры ----------


class Transaction(BaseModel):
    date: str
    amount: Decimal
    type: str
    category: Optional[str] = None
    description: Optional[str] = None


class ForecastRequest(BaseModel):
    granularity: str
    steps: int
    model: Optional[str] = None
    transactions: List[Transaction]


class ForecastResponse(BaseModel):
    period_end: List[str]
    income_forecast: List[float]
    expense_forecast: List[float]


class AdviceRequest(BaseModel):
    question: Optional[str] = None
    transactions: List[Transaction] = []


class AdviceResponse(BaseModel):
    advice: str


class ReceiptResponse(BaseModel):
    total: Optional[float]


# ---------- Стартуем и заранее грузим модели ----------

advice_tokenizer = None
advice_model = None


@app.on_event("startup")
def load_models():
    """
    Загружаем Qwen один раз при старте сервиса.
    Donut для чеков грузится в receipt_total_api при первом импорте.
    """
    global advice_tokenizer, advice_model
    advice_tokenizer, advice_model = advice_load()


# ---------- Эндпоинты ----------


from datetime import datetime  # вверху файла, если ещё не импортирован

@app.post("/forecast", response_model=ForecastResponse)
def forecast(req: ForecastRequest):
    if not req.transactions:
        raise HTTPException(status_code=400, detail="transactions is empty")

    df = pd.DataFrame([t.dict() for t in req.transactions])

    gran = (req.granularity or "month").lower()
    freq = "A-DEC" if gran.startswith("y") else "M"
    steps = int(req.steps or 1)
    method = (req.model or "auto").lower()

    inc, exp, _ = prepare_components_series(df, freq=freq)
    inc_fc = fit_and_forecast(inc, steps, freq, method=method)
    exp_fc = fit_and_forecast(exp, steps, freq, method=method)

    # Преобразуем индексы в RFC3339: "YYYY-MM-DDT00:00:00Z"
    period_end = [
        d.strftime("%Y-%m-%dT%H:%M:%SZ")
        for d in inc_fc.index.to_pydatetime().tolist()
    ]

    return ForecastResponse(
        period_end=period_end,
        income_forecast=[float(x) for x in inc_fc.values.tolist()],
        expense_forecast=[float(x) for x in exp_fc.values.tolist()],
    )


@app.post("/advice", response_model=AdviceResponse)
def advice(req: AdviceRequest):
    tx = [t.dict() for t in req.transactions] if req.transactions else []
    df = pd.DataFrame(tx) if tx else None
    snap = current_month_snapshot(df) if df is not None and not df.empty else {}

    if snap:
        ctx = [
            f"Месяц: {snap['month']}",
            f"Доход: {snap['income_total']:.0f}",
            f"Расход: {abs(snap['expense_total']):.0f}",
            f"Нетто: {snap['net']:.0f}",
        ]
        if snap.get("top_expense_categories"):
            ctx.append("Топ статей расходов:")
            for cat, val in snap["top_expense_categories"]:
                ctx.append(f"- {cat}: {abs(val):.0f}")
        context = "\n".join(ctx)
    else:
        context = "Данных за текущий месяц нет."

    question = (req.question or "").strip()

    system_msg = (
        "Ты финансовый помощник. Отвечай по-русски. "
        "Верни ТОЛЬКО список из 5–7 конкретных шагов экономии с цифрами (лимиты, проценты, частота). "
        "Каждая строка должна начинаться с символов \"- \". Никаких вступлений."
    )
    messages = [
        {"role": "system", "content": system_msg},
        {
            "role": "user",
            "content": (
                f"Мои данные за текущий месяц:\n{context}\n\nВопрос: {question}\n"
                "Начни ответ сразу со строки, которая начинается с \"- \". Верни только список."
            ),
        },
    ]

    # 1) детерминированный прогон
    raw = advice_gen(messages, advice_tokenizer, advice_model, det=True)
    text = _to_bullets(clean_ru(raw))

    # 2) если мало пунктов — пробуем стохастический прогон
    if text.count("\n") + 1 < 3:
        raw2 = advice_gen(messages, advice_tokenizer, advice_model, det=False)
        text2 = _to_bullets(clean_ru(raw2))
        if text2:
            text = text2

    # 3) Fallback: если всё равно пусто или там явный мусор — жёстко подставляем шаблон
    bad = not text.strip() or "Ты финансовый помощник" in text or "Мои данные за текущий месяц" in text

    if bad:
        # можно потом сделать это умнее и завязать на snap, но пока — вменяемый дефолт
        text = (
            "• Ограничьте траты на еду до 15–20% от дохода (для вас ≈ 15 000 ₽ в месяц) и ведите учёт покупок в одном приложении.\n"
            "• Введите недельный лимит на продукты (например, 3 500–4 000 ₽) и не выходите за него, планируя меню заранее.\n"
            "• Пересмотрите все платные подписки раз в месяц и отмените те, которыми не пользовались 30 дней и более.\n"
            "• Сгруппируйте подписки по важности: обязательные, полезные, «для развлечения» — и урежьте последнюю группу минимум на 50%.\n"
            "• Переключите часть подписок на семейные/годовые тарифы и делите стоимость с друзьями или партнёром.\n"
            "• Перенесите оплату подписок на один день месяца и установите напоминание за 3 дня до списания, чтобы успевать отменять лишнее.\n"
            "• Освободившиеся деньги (5–10% дохода) автоматически переводите на отдельный накопительный счёт сразу после зарплаты."
        )

    return AdviceResponse(advice=text)



@app.post("/receipt-total-file", response_model=ReceiptResponse)
async def receipt_total_file(file: UploadFile = File(...)):
    """
    Получает файл чека (multipart/form-data, field "file"),
    сохраняет во временный файл, вызывает extract_total и возвращает сумму.
    """
    # сохраняем во временный файл
    suffix = os.path.splitext(file.filename or "")[1] or ".jpg"
    with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp:
        contents = await file.read()
        tmp.write(contents)
        tmp_path = tmp.name

    try:
        total = extract_total(tmp_path)
        return ReceiptResponse(total=total)
    finally:
        try:
            os.remove(tmp_path)
        except OSError:
            pass


# Для локального запуска, на HF Space это не используется
if __name__ == "__main__":
    import uvicorn

    uvicorn.run("app:app", host="0.0.0.0", port=7860, workers=1)