from fastapi import APIRouter, Query, HTTPException
from sqlalchemy.orm import Session
from typing import List, Dict
from datetime import datetime, timedelta
from models.market import MarketData, MarketDataDB
from db import SessionLocal

router = APIRouter()

# Only allow these markets
SUPPORTED_MARKETS: Dict[str, str] = {
    "ISONE": "ISONE",
    "MISO": "MISO",
    "NYISO": "NYISO",
}

def check_market_supported(market: str):
    market = market.upper()
    if market not in SUPPORTED_MARKETS:
        raise HTTPException(status_code=400, detail=f"Unsupported market '{market}'. Supported: {list(SUPPORTED_MARKETS.keys())}")
    return market

@router.get("/day-ahead", response_model=List[MarketData])
def get_day_ahead_data(market: str = Query("ISONE")):
    db: Session = SessionLocal()
    market = check_market_supported(market)

    records = db.query(MarketDataDB)\
        .filter(MarketDataDB.market == market, MarketDataDB.type == "DAYAHEAD")\
        .order_by(MarketDataDB.timestamp)\
        .all()

    db.close()
    return [
        MarketData(
            timestamp=r.timestamp,
            lmp=r.lmp,
            energy=r.energy,
            congestion=r.congestion,
            loss=r.loss,
        ) for r in records
    ]

@router.get("/real-time", response_model=List[MarketData])
def get_real_time_data(market: str = Query("ISONE")):
    db: Session = SessionLocal()
    market = check_market_supported(market)

    start_time = datetime.utcnow() - timedelta(days=1)
    records = db.query(MarketDataDB)\
        .filter(
            MarketDataDB.market == market,
            MarketDataDB.type == "REALTIME",
            MarketDataDB.timestamp >= start_time
        )\
        .order_by(MarketDataDB.timestamp)\
        .all()

    db.close()
    return [
        MarketData(
            timestamp=r.timestamp,
            lmp=r.lmp,
            energy=r.energy,
            congestion=r.congestion,
            loss=r.loss,
        ) for r in records
    ]