import pandas as pd
import requests
import xml.etree.ElementTree as ET
from datetime import datetime
import time
import json
import os
from pathlib import Path


class OECDDataPuller:
    """
    OECD data puller that saves progress and retries until all indicators are fetched.
    Designed to handle API rate limits by saving state between runs.
    """

    INDICATOR_CONFIG = [
        {
            "name": "Business Confidence Index",
            "series": "BCICP",
            "dataset": "SDD.STES,DSD_STES@DF_CLI,",
            "filter": ".....",
            "col_name": "macro_business_confidence_index",
        },
        {
            "name": "Consumer Confidence Index",
            "series": "CCICP",
            "dataset": "SDD.STES,DSD_STES@DF_CLI,",
            "filter": ".....",
            "col_name": "macro_consumer_confidence_index",
        },
        {
            "name": "CPI Total",
            "series": "N.CPI",
            "dataset": "SDD.TPS,DSD_PRICES@DF_PRICES_ALL,",
            "filter": "PA._T.N.GY",
            "col_name": "macro_cpi_total",
        },
        {
            "name": "CPI Housing",
            "series": "N.CPI",
            "dataset": "SDD.TPS,DSD_PRICES@DF_PRICES_ALL,",
            "filter": "PA.CP041T043.N.GY",
            "col_name": "macro_cpi_housing",
        },
        {
            "name": "CPI Food",
            "series": "N.CPI",
            "dataset": "SDD.TPS,DSD_PRICES@DF_PRICES_ALL,",
            "filter": "PA.CP01.N.GY",
            "col_name": "macro_cpi_food",
        },
        {
            "name": "CPI Energy",
            "series": "N.CPI",
            "dataset": "SDD.TPS,DSD_PRICES@DF_PRICES_ALL,",
            "filter": "PA.CP045_0722.N.GY",
            "col_name": "macro_cpi_energy",
        },
        {
            "name": "Unemployment Rate",
            "series": "UNE_LF_M",
            "dataset": "SDD.TPS,DSD_LFS@DF_IALFS_UNE_M,",
            "filter": "._Z.Y._T.Y_GE15.",
            "col_name": "macro_unemployment_rate",
            "special": "SPECIAL_UNE",
        },
        {
            "name": "Real House Prices",
            "series": "RHP",
            "dataset": "ECO.MPD,DSD_AN_HOUSE_PRICES@DF_HOUSE_PRICES,1.0",
            "filter": "",
            "col_name": "macro_real_house_prices",
        },
        {
            "name": "Manufacturing Production",
            "series": "PRVM",
            "dataset": "SDD.STES,DSD_KEI@DF_KEI,4.0",
            "filter": "IX.C..",
            "col_name": "macro_manufacturing_production_volume",
        },
        {
            "name": "Retail Trade Volume",
            "series": "TOVM",
            "dataset": "SDD.STES,DSD_KEI@DF_KEI,4.0",
            "filter": "IX...",
            "col_name": "macro_retail_trade_volume",
        },
        {
            "name": "Interbank Rate",
            "series": "IRSTCI",
            "dataset": "SDD.STES,DSD_KEI@DF_KEI,4.0",
            "filter": "PA...",
            "col_name": "macro_interbank_rate",
        },
        {
            "name": "Long-term Interest Rate",
            "series": "IRLT",
            "dataset": "SDD.STES,DSD_KEI@DF_KEI,4.0",
            "filter": "PA...",
            "col_name": "macro_long_term_interest_rate",
        },
        {
            "name": "GDP Growth",
            "series": "B1GQ",
            "dataset": "SDD.NAD,DSD_NAMAIN1@DF_QNA,1.1",
            "filter": "._Z....GY.T0102",
            "col_name": "macro_gdp_growth_yoy",
            "special": "SPECIAL_GDP",
        },
    ]

    def __init__(self, country="GBR", start_date="2020-01-01", output_dir=None):
        """
        Initialize the puller.

        Args:
            country (str): Country code (e.g., "GBR")
            start_date (str): Start date for data collection
            output_dir (str): Directory to save output files and state. 
                           Defaults to shared network path if available, else local "oecd_data"
        """
        self.country = country
        self.start_date = start_date
        
        # Determine output directory: try shared path first, fall back to local
        if output_dir is None:
            user_home = os.path.expanduser("~")
            shared_path = Path(user_home) / "im-sciences.com" / "FileShare - MasterDrive" / "Central Database" / "Pull All" / "OECD Database"
            local_path = Path("oecd_data")
            
            # Try to use shared path if it exists and is accessible
            if shared_path.exists() and shared_path.is_dir():
                self.output_dir = shared_path
                print(f"Using shared network path: {self.output_dir}")
            else:
                self.output_dir = local_path
                print(f"Shared path not available. Using local directory: {self.output_dir}")
        else:
            self.output_dir = Path(output_dir)
        
        self.output_dir.mkdir(parents=True, exist_ok=True)

        self.state_file = self.output_dir / f"state_{country}.json"
        self.data_file = self.output_dir / f"oecd_data_{country}.csv"
        self.log_file = self.output_dir / f"log_{country}.txt"

        self.state = self._load_state()

    def _load_state(self):
        """Load the current state from file, or initialize a new state."""
        if self.state_file.exists():
            with open(self.state_file, 'r') as f:
                return json.load(f)
        else:
            return {
                "completed_indicators": [],
                "failed_attempts": {},
                "last_run": None,
                "fully_complete": False
            }

    def _save_state(self):
        """Save the current state to file."""
        self.state["last_run"] = datetime.now().isoformat()
        with open(self.state_file, 'w') as f:
            json.dump(self.state, f, indent=2)

    def _log(self, message):
        """Write a log message to both console and log file."""
        timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        log_message = f"[{timestamp}] {message}"
        print(log_message)

        with open(self.log_file, 'a', encoding='utf-8') as f:
            f.write(log_message + "\n")

    @staticmethod
    def parse_quarter(date_str):
        """Parse a string in 'YYYY-Q#' format into a datetime object."""
        year, quarter = date_str.split("-")
        quarter_number = int(quarter[1])
        month = (quarter_number - 1) * 3 + 1
        return pd.Timestamp(f"{year}-{month:02d}-01")

    def _build_url(self, series, dataset_id, filter_val, freq, special_flag=None):
        """Build the appropriate OECD API URL based on indicator type."""
        if special_flag == "SPECIAL_GDP":
            return f"https://sdmx.oecd.org/public/rest/data/OECD.{dataset_id}/{freq}..{self.country}...{series}.{filter_val}?startPeriod=1950-01"
        elif special_flag == "SPECIAL_UNE":
            return f"https://sdmx.oecd.org/public/rest/data/OECD.{dataset_id}/{self.country}.{series}.{filter_val}.{freq}?startPeriod=1950-01"
        else:
            return f"https://sdmx.oecd.org/public/rest/data/OECD.{dataset_id}/{self.country}.{freq}.{series}.{filter_val}?startPeriod=1950-01"

    def _extract_observations(self, xml_content):
        """Extract dates and values from OECD API XML response."""
        root = ET.fromstring(xml_content)
        namespaces = {
            "generic": "http://www.sdmx.org/resources/sdmxml/schemas/v2_1/data/generic",
        }

        dates = []
        values = []

        for obs in root.findall(".//generic:Obs", namespaces):
            time_period = obs.find(".//generic:ObsDimension", namespaces).get("value")
            value = obs.find(".//generic:ObsValue", namespaces).get("value")

            if time_period and value:
                dates.append(time_period)
                values.append(float(value))

        return dates, values

    def _fetch_indicator_data(self, series, dataset_id, filter_val, col_name, special_flag=None):
        """
        Attempt to fetch data for a single indicator across different frequencies.

        Returns:
            tuple: (DataFrame or None, frequency_used, success_flag)
        """
        for freq in ["M", "Q", "A"]:
            url = self._build_url(series, dataset_id, filter_val, freq, special_flag)

            try:
                response = requests.get(url, timeout=15)

                if response.status_code == 429:
                    self._log(f"Rate limit hit for {col_name}")
                    return None, None, False

                if response.status_code != 200:
                    continue

                dates, values = self._extract_observations(response.content)

                if len(dates) == 0:
                    continue

                data = pd.DataFrame({"OBS": dates, col_name: values})

                if freq == "Q":
                    data["OBS"] = data["OBS"].apply(self.parse_quarter)
                else:
                    data["OBS"] = data["OBS"].apply(lambda x: datetime.strptime(x, "%Y-%m"))

                data.sort_values(by="OBS", inplace=True)

                return data, freq, True

            except Exception as e:
                self._log(f"Error fetching {col_name}: {str(e)}")
                continue

        return None, None, False

    def _load_existing_data(self):
        """Load existing data from CSV file if it exists."""
        if self.data_file.exists():
            df = pd.read_csv(self.data_file)
            df['OBS'] = pd.to_datetime(df['OBS'])
            return df
        else:
            date_range = pd.date_range(start=self.start_date, end=datetime.today(), freq="D")
            return pd.DataFrame({"OBS": date_range})

    def _save_data(self, df):
        """Save DataFrame to CSV file."""
        df.to_csv(self.data_file, index=False)

    def fetch_pending_indicators(self):
        """Fetch all indicators to ensure the most up-to-date data."""
        # Reset completed indicators for this fresh run
        self.state["completed_indicators"] = []
        
        daily_df = self._load_existing_data()

        # Always attempt to refresh all indicators
        pending_indicators = self.INDICATOR_CONFIG

        self._log(f"Starting fetch cycle. Pending: {len(pending_indicators)}/{len(self.INDICATOR_CONFIG)}")

        for indicator in pending_indicators:
            col_name = indicator["col_name"]

            self._log(f"Fetching: {indicator['name']} ({col_name})")

            data, freq_used, success = self._fetch_indicator_data(
                indicator["series"],
                indicator["dataset"],
                indicator["filter"],
                col_name,
                indicator.get("special")
            )

            if success:
                self._log(f"  [OK] Success! {len(data)} observations ({freq_used})")

                daily_df = pd.merge_asof(
                    daily_df,
                    data[["OBS", col_name]],
                    on="OBS",
                    direction="backward",
                )

                self.state["completed_indicators"].append(col_name)

                if col_name in self.state["failed_attempts"]:
                    del self.state["failed_attempts"][col_name]

                self._save_data(daily_df)
                self._save_state()

                time.sleep(2)

            else:
                self._log(f"  [FAIL] Failed to fetch {col_name}")

                if col_name not in self.state["failed_attempts"]:
                    self.state["failed_attempts"][col_name] = 0
                self.state["failed_attempts"][col_name] += 1

                self._save_state()

        remaining = len(self.INDICATOR_CONFIG) - len(self.state["completed_indicators"])

        if remaining == 0:
            self._log("All indicators successfully fetched!")
        else:
            self._log(f"Fetch cycle complete. {remaining} indicators still have failures.")
        
        return True  # Always return True to indicate the refresh cycle completed

    def run_until_complete(self, max_iterations=None, sleep_seconds=3600):
        """
        Run the fetcher to refresh all indicators with up-to-date data.

        Args:
            max_iterations (int): Maximum number of cycles to run (None = unlimited)
            sleep_seconds (int): Seconds to wait between cycles (default: 3600 = 1 hour)
        """
        iteration = 0

        self._log("=" * 80)
        self._log(f"Starting OECD data puller for {self.country}")
        self._log(f"Output directory: {self.output_dir.absolute()}")
        self._log("=" * 80)

        while True:
            iteration += 1

            if max_iterations and iteration > max_iterations:
                self._log(f"Reached maximum iterations ({max_iterations}). Stopping.")
                break

            self._log(f"{'='*80}")
            self._log(f"ITERATION {iteration}")
            self._log(f"{'='*80}")

            self.fetch_pending_indicators()

            # Check if all indicators have data
            if len(self.state["completed_indicators"]) == len(self.INDICATOR_CONFIG):
                self._log("=" * 80)
                self._log("SUCCESS! All indicators have been fetched.")
                self._log(f"Data saved to: {self.data_file.absolute()}")
                self._log("=" * 80)
                break

            self._log(f"Waiting {sleep_seconds} seconds before next attempt...")
            self._log(f"Next run scheduled for: {(datetime.now() + pd.Timedelta(seconds=sleep_seconds)).strftime('%Y-%m-%d %H:%M:%S')}")

            time.sleep(sleep_seconds)

    def get_status(self):
        """Get current status of data collection."""
        total = len(self.INDICATOR_CONFIG)
        completed = len(self.state["completed_indicators"])

        status = {
            "country": self.country,
            "total_indicators": total,
            "completed_indicators": completed,
            "remaining_indicators": total - completed,
            "completion_percentage": (completed / total) * 100,
            "fully_complete": self.state["fully_complete"],
            "last_run": self.state["last_run"],
            "failed_attempts": self.state["failed_attempts"],
            "completed_list": self.state["completed_indicators"]
        }

        return status


def main():
    """Main execution function with example usage."""

    # Example 1: Run until complete with 1-hour intervals
    puller = OECDDataPuller(
        country="GBR",
        start_date="2020-01-01",
        output_dir="oecd_data"
    )

    # Check current status
    status = puller.get_status()
    print(f"Current progress: {status['completed_indicators']}/{status['total_indicators']} "
          f"({status['completion_percentage']:.1f}%)")

    # Run until complete (will retry every hour)
    puller.run_until_complete(sleep_seconds=3600)

    # Example 2: Run a single fetch cycle (useful for manual/scheduled execution)
    # puller = OECDDataPuller(country="GBR")
    # puller.fetch_pending_indicators()

    # Example 3: Run with a custom sleep interval (e.g., 30 minutes)
    # puller = OECDDataPuller(country="USA")
    # puller.run_until_complete(sleep_seconds=1800)


if __name__ == "__main__":
    main()
