from __future__ import annotations

import warnings
from datetime import date
from dateutil.relativedelta import relativedelta

import numpy as np
import pandas as pd

warnings.filterwarnings("ignore")


def _project_months(from_date: date, n: int) -> list[str]:
    months = []
    d = from_date.replace(day=1)
    for i in range(1, n + 1):
        months.append((d + relativedelta(months=i)).strftime("%Y-%m"))
    return months


def project_net_worth(df: pd.DataFrame, years: int = 5) -> dict:
    """
    df columns: ds (monthly datetime), y (net_worth float)
    Returns history + 3-scenario projections.
    """
    n_months = years * 12
    today = date.today()
    future_dates = _project_months(today, n_months)

    history = [
        {"date": row["ds"].strftime("%Y-%m"), "value": round(float(row["y"]), 2)}
        for _, row in df.iterrows()
    ]

    if df.empty or len(df) < 2:
        # No data — return flat projection from 0
        last_val = float(df["y"].iloc[-1]) if not df.empty else 0.0
        flat = [{"date": d, "value": round(last_val, 2)} for d in future_dates]
        return {
            "history": history,
            "projections": {"conservative": flat, "base": flat, "optimistic": flat},
            "insufficient_data": True,
        }

    try:
        from statsmodels.tsa.holtwinters import ExponentialSmoothing

        values = df["y"].tolist()

        if len(values) >= 12:
            model = ExponentialSmoothing(values, trend="add", seasonal="add", seasonal_periods=12)
        elif len(values) >= 4:
            model = ExponentialSmoothing(values, trend="add", seasonal=None)
        else:
            model = ExponentialSmoothing(values, trend="add", seasonal=None)

        fit = model.fit(optimized=True, disp=False)
        base_fcast = fit.forecast(n_months)

        # Estimate monthly trend from the fit
        monthly_trend = float(np.mean(np.diff(base_fcast[:12]))) if len(base_fcast) >= 12 else 0.0
        last_val = float(values[-1])

        # Scale trends for scenarios
        def build_scenario(scale: float) -> list[dict]:
            pts = []
            v = last_val
            for i, d in enumerate(future_dates):
                v = float(base_fcast[i]) + (scale - 1.0) * monthly_trend * (i + 1)
                pts.append({"date": d, "value": round(v, 2)})
            return pts

        return {
            "history": history,
            "projections": {
                "conservative": build_scenario(0.5),
                "base": [{"date": d, "value": round(float(v), 2)} for d, v in zip(future_dates, base_fcast)],
                "optimistic": build_scenario(1.5),
            },
            "insufficient_data": False,
        }

    except Exception:
        # Fallback: linear trend from last 2 values
        trend = float(df["y"].iloc[-1]) - float(df["y"].iloc[-2])
        last_val = float(df["y"].iloc[-1])

        def linear_scenario(t_scale: float) -> list[dict]:
            return [
                {"date": d, "value": round(last_val + t_scale * trend * (i + 1), 2)}
                for i, d in enumerate(future_dates)
            ]

        return {
            "history": history,
            "projections": {
                "conservative": linear_scenario(0.5),
                "base": linear_scenario(1.0),
                "optimistic": linear_scenario(1.5),
            },
            "insufficient_data": False,
        }