# app.py
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from typing import List, Literal
import joblib
import numpy as np

app = FastAPI(title="Transaction Categorizer")

# Load models once at startup
shop_model     = joblib.load("shop_classifier.pkl")
category_model = joblib.load("category_classifier.pkl")

def predict_or_na(pipeline, text: str, threshold: float = 0.3) -> str:
    """
    Returns the top class if its probability ≥ threshold, else "N/A"
    """
    probs = pipeline.predict_proba([text])[0]
    top_idx = int(np.argmax(probs))
    return str(pipeline.classes_[top_idx]) if probs[top_idx] >= threshold else "N/A"

class TransactionIn(BaseModel):
    id: str
    description: str

class TransactionOut(BaseModel):
    id: str
    description: str
    shop: str
    category: str

@app.post("/predict", response_model=List[TransactionOut])
def predict_transactions(
    items: List[TransactionIn],
    threshold: float = 0.3
):
    """
    Predict shop and category for each transaction.
    - **items**: list of `{ id, description }`
    - **threshold**: optional override for probability cutoff
    """
    results = []
    for item in items:
        desc_norm = item.description.lower().strip()
        shop_pred     = predict_or_na(shop_model,     desc_norm, threshold)
        category_pred = predict_or_na(category_model, desc_norm, threshold)
        results.append(
            TransactionOut(
                id=item.id,
                description=item.description,
                shop=shop_pred,
                category=category_pred
            )
        )
    return results

# Optional healthcheck
@app.get("/health")
def health():
    return {"status": "ok"}