from sqlalchemy import select, extract, func
from typing import Dict
from decimal import Decimal
from collections import defaultdict

from src.financial.model import GeneralLedger
from .models import  ForecastData  # ForecastData must exist
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.dialects.postgresql import insert 

async def get_quarterly_revenue(year: int, db: AsyncSession) -> Dict[int, float]:
    stmt = (
        select(
            extract('quarter', GeneralLedger.date).label('q'),
            func.sum(GeneralLedger.credit).label('amt')
        )
        .filter(
            extract('year', GeneralLedger.date) == year,
            GeneralLedger.account_name.ilike("%revenue%")
        )
        .group_by('q')
    )
    result = await db.execute(stmt)
    rows = result.all()
    revenue_by_quarter = {int(row.q): float(row.amt or 0) for row in rows}
    for q in range(1, 5):
        revenue_by_quarter.setdefault(q, 0.0)
    return revenue_by_quarter


def calculate_scenarios(baseline: Dict[int, float]) -> Dict[str, Dict[int, float]]:
    best_case = {q: val * 1.15 for q, val in baseline.items()}
    worst_case = {q: val * 0.85 for q, val in baseline.items()}
    return {
        "Baseline": baseline,
        "Best Case": best_case,
        "Worst Case": worst_case,
    }

from sqlalchemy.dialects.postgresql import insert  # Import PostgreSQL-specific insert
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession
from typing import Dict
async def save_forecast(
    scenario: str,
    year: int,
    data: Dict[int, float],
    is_ai: bool,
    db: AsyncSession,
    metric_type: str  # must be passed in now
):
    for quarter, value in data.items():
        stmt = insert(ForecastData).values(
            year=year,
            quarter=quarter,
            scenario=scenario,
            is_ai=is_ai,
            value=value,
            metric_type=metric_type
        ).on_conflict_do_update(
            index_elements=[ForecastData.year, ForecastData.quarter, ForecastData.scenario, ForecastData.is_ai],
            set_={"value": value}
        ).returning(ForecastData.id)

        await db.execute(stmt)
    await db.commit()

async def get_saved_forecast(
    year: int,
    scenario_name: str,
    db: AsyncSession
) -> Dict[str, Dict[int, float]]:
    stmt = select(
        ForecastData.quarter,
        ForecastData.value,
        ForecastData.is_ai
    ).filter(
        ForecastData.year == year,
        ForecastData.scenario == scenario_name
    )
    result = await db.execute(stmt)
    rows = result.all()

    ai_data = {}
    user_data = {}
    for q, val, is_ai in rows:
        if is_ai:
            ai_data[q] = val
        else:
            user_data[q] = val

    return {
        "AI Suggested": ai_data,
        "User Forecast": user_data
    }



