import logging
import os
import time
from pathlib import Path
import httpx
import pandas as pd

logging.basicConfig(
    level=logging.INFO, 
    format="[%(asctime)s] %(message)s",
    datefmt="%H:%M:%S",
    force=True 
)
logging.getLogger("httpx").setLevel(logging.WARNING)
log = logging.getLogger(__name__)

class SmardFetcher:
    BASE_URL = "https://www.smard.de/app/chart_data"
    RESOLUTION = "hour"    
    START_DATE = pd.Timestamp("2020-01-01", tz="UTC")
    
    def __init__(self, data_dir: str = None):
        target_dir = data_dir or os.getenv("DATA_DIR", "data")
        self.raw_dir = Path(target_dir) / "raw"
        self.raw_dir.mkdir(parents=True, exist_ok=True)
        self.dataset_file = self.raw_dir / "pvar_dataset.csv"

    def _get_relevant_timestamps(self, client: httpx.Client, filter_id: str, region_id: str, last_dt: pd.Timestamp = None) -> list:
        url = f"{self.BASE_URL}/{filter_id}/{region_id}/index_{self.RESOLUTION}.json"
        response = client.get(url)
        if response.status_code != 200:
            return []
            
        timestamps = response.json().get("timestamps", [])
        start_ms = self.START_DATE.timestamp() * 1000
        valid_ts = [ts for ts in timestamps if ts >= start_ms]
        
        if last_dt:
            last_ms = last_dt.timestamp() * 1000
            valid_ts = [ts for ts in valid_ts if ts > last_ms]
            
        return valid_ts

    def _fetch_series(self, client: httpx.Client, filter_id: str, region_id: str, timestamp: int, col_name: str, retries=3) -> pd.DataFrame:
        url = f"{self.BASE_URL}/{filter_id}/{region_id}/{filter_id}_{region_id}_{self.RESOLUTION}_{timestamp}.json"
        for attempt in range(retries):
            response = client.get(url)
            if response.status_code == 200:
                data = response.json().get("series", [])
                if not data:
                    return pd.DataFrame()
                df = pd.DataFrame(data, columns=["timestamp", col_name])
                df["timestamp"] = pd.to_datetime(df["timestamp"], unit="ms", utc=True)
                return df.set_index("timestamp")
            elif response.status_code == 429:
                time.sleep(2)
            else:
                break
        return pd.DataFrame()

    def _sync_feature(self, client: httpx.Client, filter_id: str, region_id: str, col_name: str, last_dt: pd.Timestamp = None) -> pd.DataFrame:
        timestamps = self._get_relevant_timestamps(client, filter_id, region_id, last_dt)
        if not timestamps:
            return pd.DataFrame()

        log.info(f"  -> {col_name:<18} ({region_id}) | Downloading {len(timestamps)} blocks...")
        dfs = []
        for ts in timestamps:
            dfs.append(self._fetch_series(client, filter_id, region_id, ts, col_name))
            time.sleep(0.02) # Rate limit safety
                
        valid_dfs = [df for df in dfs if not df.empty]
        return pd.concat(valid_dfs).sort_index() if valid_dfs else pd.DataFrame()

    def sync_all(self, progress_cb=None) -> pd.DataFrame:
        # Hard wipe of the corrupted file to start 100% fresh
        if self.dataset_file.exists():
            log.warning("Wiping corrupted dataset to build a seamless combined version...")
            self.dataset_file.unlink(missing_ok=True)

        feature_dfs = []
        with httpx.Client(timeout=30.0) as client:
            # 1. Fetch Era 1 & 2 Prices
            if progress_cb:
                progress_cb("Lade Strompreise herunter...")
            df_price_legacy = self._sync_feature(client, "4169", "DE-LU", "price_mwh")
            df_price_modern = self._sync_feature(client, "4169", "DE", "price_mwh")
            
            # Combine both price eras vertically, drop overlapping duplicates
            df_price = pd.concat([df_price_legacy, df_price_modern]).sort_index()
            df_price = df_price[~df_price.index.duplicated(keep='first')]
            
            if not df_price.empty:
                feature_dfs.append(df_price)

            # 3. Fetch grid physics (Always under DE)
            grid_filters = {
                "410": "load_total",
                "125": "prog_pv",
                "123": "prog_wind_onshore",
                "3791": "prog_wind_offshore"
            }
            grid_names = {
                "load_total": "Netzlast",
                "prog_pv": "Solarprognose",
                "prog_wind_onshore": "Wind Onshore Prognose",
                "prog_wind_offshore": "Wind Offshore Prognose"
            }
            for fid, col in grid_filters.items():
                if progress_cb:
                    progress_cb(f"Lade {grid_names.get(col, col)} herunter...")
                df_feat = self._sync_feature(client, fid, "DE", col)
                if not df_feat.empty:
                    feature_dfs.append(df_feat)

        if not feature_dfs:
            log.error("No data could be retrieved.")
            return pd.DataFrame()

        if progress_cb:
            progress_cb("Verarbeite Daten...")
        # Merge all features side-by-side on the timestamp index
        df_final = pd.concat(feature_dfs, axis=1).sort_index()
        
        # Keep only data from 2020 onwards
        df_final = df_final[df_final.index >= self.START_DATE]
        
        df_final.to_csv(self.dataset_file)
        log.info(f"Successfully generated clean dataset: {self.dataset_file} ({len(df_final)} rows)")
        return df_final

if __name__ == "__main__":
    fetcher = SmardFetcher()
    fetcher.sync_all()