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_info

app = FastAPI()

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():
    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)

    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"
                "Начни ответ сразу со строки, которая начинается с \"- \". Верни только список."
            ),
        },
    ]

    raw = advice_gen(messages, advice_tokenizer, advice_model, det=True)
    text = _to_bullets(clean_ru(raw))

    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

    return AdviceResponse(advice=text)



@app.post("/receipt-total-file", response_model=ReceiptResponse)
async def receipt_total_file(file: UploadFile = File(...)):

    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

if __name__ == "__main__":
    import uvicorn

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