from datetime import datetime, timedelta, UTC
from sqlalchemy.orm import Session
from db import SessionLocal
from models.market import MarketDataDB
from gridstatus import ISONE, MISO, NYISO
from typing import Dict, Type

MARKET_CLASSES: Dict[str, Type] = {
    "ISONE": ISONE,
    "MISO": MISO,
    "NYISO": NYISO,
}

def get_iso_instance(market: str):
    market = market.upper()
    if market not in MARKET_CLASSES:
        raise ValueError(f"Unsupported market '{market}'. Supported: {list(MARKET_CLASSES.keys())}")
    return MARKET_CLASSES[market]()

def update_market_data():
    db: Session = SessionLocal()

    for market_name in MARKET_CLASSES.keys():
        print(f"Processing {market_name}")
        iso = get_iso_instance(market_name)

        # --- Real-Time Data (5 min) ---
        last_realtime = db.query(MarketDataDB)\
            .filter(MarketDataDB.market == market_name, MarketDataDB.type == "REALTIME")\
            .order_by(MarketDataDB.timestamp.desc())\
            .first()

        if not last_realtime or (datetime.now(UTC) - last_realtime.timestamp.replace(tzinfo=UTC) > timedelta(minutes=5)):
            print(f"Getting realtime data for {market_name}")
            df = iso.get_lmp(date="latest", market="REAL_TIME_5_MIN", locations="ALL")
            df["Interval Start"] = df["Interval Start"].dt.tz_convert('UTC')
            grouped = df.groupby("Interval Start")[["LMP", "Energy", "Congestion", "Loss"]].mean().reset_index()

            for _, row in grouped.iterrows():
                if last_realtime and row["Interval Start"] <= last_realtime.timestamp.replace(tzinfo=UTC):
                    continue  # Skip old data
                entry = MarketDataDB(
                    timestamp=row["Interval Start"],
                    lmp=row["LMP"],
                    energy=row["Energy"],
                    congestion=row["Congestion"],
                    loss=row["Loss"],
                    market=market_name,
                    type="REALTIME",
                )
                db.add(entry)

        # --- Day-Ahead Hourly Data (1 hour) ---
        last_dayahead = db.query(MarketDataDB)\
            .filter(MarketDataDB.market == market_name, MarketDataDB.type == "DAYAHEAD")\
            .order_by(MarketDataDB.timestamp.desc())\
            .first()

        if not last_dayahead or (datetime.now(UTC) - last_dayahead.timestamp.replace(tzinfo=UTC) > timedelta(hours=1)):
            print(f"Getting day-ahead data for {market_name}")
            now_utc = datetime.now(UTC)
            day_ahead_date = now_utc.date()
            if now_utc.hour >= 18:  # After 6PM UTC, markets usually publish next day's data
                day_ahead_date += timedelta(days=1)
            df = iso.get_lmp(date=day_ahead_date, market="DAY_AHEAD_HOURLY", locations="ALL")
            df["Interval Start"] = df["Interval Start"].dt.tz_convert('UTC')
            grouped = df.groupby("Interval Start")[["LMP", "Energy", "Congestion", "Loss"]].mean().reset_index()

            for _, row in grouped.iterrows():
                if last_dayahead and row["Interval Start"] <= last_dayahead.timestamp.replace(tzinfo=UTC):
                    continue
                entry = MarketDataDB(
                    timestamp=row["Interval Start"],
                    lmp=row["LMP"],
                    energy=row["Energy"],
                    congestion=row["Congestion"],
                    loss=row["Loss"],
                    market=market_name,
                    type="DAYAHEAD",
                )
                db.add(entry)

    db.commit()
    db.close()

if __name__ == "__main__":
    update_market_data()