from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import pandas as pd
from prophet import Prophet
import psycopg2
import os
app = FastAPI()

db = os.getenv('db')

# Database URL
DATABASE_URL = db

# Request body schema
class PredictionRequest(BaseModel):
    user_id: str
    tanggal_prediksi: str  # Format YYYY-MM-DD

# Fungsi ambil data transaksi
def get_transactions(user_id):
    try:
        conn = psycopg2.connect(DATABASE_URL)
        cur = conn.cursor()
        cur.execute("""
            SELECT id, user_id, amount, type, description, transaction_date 
            FROM transactions 
            WHERE user_id = %s
            ORDER BY transaction_date
        """, (user_id,))
        rows = cur.fetchall()
        columns = ['id', 'user_id', 'amount', 'type', 'description', 'transaction_date']
        df = pd.DataFrame(rows, columns=columns)
        cur.close()
        conn.close()
        return df
    except Exception as e:
        return None

# Hitung grand total harian
def prepare_data(df):
    df['transaction_date'] = pd.to_datetime(df['transaction_date'])
    df = df.sort_values(by=['transaction_date', 'id'])

    grand_total = 0
    grand_totals = []
    for _, row in df.iterrows():
        if row['type'] == 'pemasukan':
            grand_total += row['amount']
        elif row['type'] == 'pengeluaran':
            grand_total -= row['amount']
        grand_totals.append(grand_total)

    df['grand_total'] = grand_totals
    return df

@app.post("/predict")
def predict_saldo(request: PredictionRequest):
    try:
        future_date = pd.to_datetime(request.tanggal_prediksi)
    except:
        raise HTTPException(status_code=400, detail="Format tanggal tidak valid. Gunakan YYYY-MM-DD.")

    df = get_transactions(request.user_id)
    if df is None or df.empty:
        raise HTTPException(status_code=404, detail="Data transaksi tidak ditemukan untuk user_id tersebut.")

    df = prepare_data(df)
    df_prophet = df[['transaction_date', 'grand_total']].rename(
        columns={'transaction_date': 'ds', 'grand_total': 'y'}
    )

    model = Prophet()
    model.fit(df_prophet)

    last_date = df_prophet['ds'].max()
    days_ahead = (future_date - last_date).days

    if days_ahead <= 0:
        raise HTTPException(status_code=400, detail=f"Tanggal prediksi ({future_date.date()}) harus setelah data terakhir ({last_date.date()}).")

    future = model.make_future_dataframe(periods=days_ahead, freq='D')
    forecast = model.predict(future)
    prediction = forecast[forecast['ds'] == future_date]

    if prediction.empty:
        raise HTTPException(status_code=404, detail="Tanggal tidak ditemukan dalam hasil prediksi.")

    result = prediction[['ds', 'yhat', 'yhat_lower', 'yhat_upper']].iloc[0]
    return {
        "tanggal": result['ds'].date().isoformat(),
        "prediksi_saldo": round(result['yhat'], 2),
        "batas_bawah": round(result['yhat_lower'], 2),
        "batas_atas": round(result['yhat_upper'], 2)
    }