import datetime
import logging

import yfinance as yf
from fredapi import Fred
import sqlalchemy as sa

from config import config
from models.base import financial_engine

fred = Fred(api_key=config.FRED_API_KEY)


def get_data_from_yf(symbol: str, start_date: str, end_date: str, all_data: bool = False):
    gold = yf.Ticker(symbol)

    if all_data:
        ticker_data = gold.history(period="max")
    else:
        ticker_data = gold.history(start=start_date, end=end_date)

    return ticker_data


def get_shsz_history_data(start_date: str, end_date: str, all_data: bool = False):
    ticker_symbol = "000300.SS"

    return get_data_from_yf(ticker_symbol, start_date, end_date, all_data)


def get_dji_history_data(start_date: str, end_date: str, all_data: bool = False):
    ticker_symbol = "^DJI"

    return get_data_from_yf(ticker_symbol, start_date, end_date, all_data)


def get_nasdaq_history_data(start_date: str, end_date: str, all_data: bool = False):
    ticker_symbol = "^IXIC"

    return get_data_from_yf(ticker_symbol, start_date, end_date, all_data)


def get_gold_history_data(start_date: str, end_date: str, all_data: bool = False):
    ticker_symbol = "GC=F"

    return get_data_from_yf(ticker_symbol, start_date, end_date, all_data)


def get_oil_history_data(start_date: str, end_date: str, all_data: bool = False):
    ticker_symbol = "CL=F"

    return get_data_from_yf(ticker_symbol, start_date, end_date, all_data)


def get_fedfunds_history_data(start_date: str, end_date: str, all_data: bool = False):
    if all_data:
        data = fred.get_series("FEDFUNDS")
    else:
        data = fred.get_series("FEDFUNDS", start_date, end_date)

    return data


def get_unrate_history_data(start_date: str, end_date: str, all_data: bool = False):
    if all_data:
        data = fred.get_series("UNRATE")
    else:
        data = fred.get_series("UNRATE", start_date, end_date)

    return data


def refresh_six_cols_history_data(table_name: str, all_data: bool = False):
    data_to_insert = []
    now = datetime.datetime.now()
    end_date = now.strftime("%Y-%m-%d")
    start_date = (now - datetime.timedelta(days=30)).strftime("%Y-%m-%d")
    fn_mapper = {
        "csi300": get_shsz_history_data,
        "gold": get_gold_history_data,
        "crude_oil": get_oil_history_data,
        "nasdaq": get_nasdaq_history_data,
        "dji": get_dji_history_data,
    }

    fn = fn_mapper.get(table_name)
    if not fn:
        logging.error("Table %s is not supported.", table_name)
        return

    with financial_engine.connect() as conn:
        for i, row in fn(start_date, end_date, all_data=all_data).iterrows():
            date = str(i).split(" ")[0]
            open = row["Open"]
            high = row["High"]
            low = row["Low"]
            close = row["Close"]
            volume = row["Volume"]

            # check if the data already exists
            if not all_data:
                result = conn.execute(sa.text(
                    f"SELECT * FROM {table_name} WHERE date = :date"
                ), {"date": date})
                if result.fetchone():
                    logging.info("Data for %s of %s already exists.", date, table_name)
                    continue

            data_to_insert.append({
                "date": date,
                "open": open,
                "high": high,
                "low": low,
                "close": close,
                "volume": volume
            })

        if data_to_insert:
            conn.execute(sa.text(
                f"INSERT INTO {table_name} (date, open, high, low, close, volume) VALUES (:date, :open, :high, :low, :close, :volume)",
            ), data_to_insert)

        conn.commit()

        logging.info("Inserted %d rows of data into %s.", len(data_to_insert), table_name)


def refresh_two_cols_history_data(table_name: str, all_data: bool = False):
    data_to_insert = []
    now = datetime.datetime.now()
    end_date = now.strftime("%Y-%m-%d")
    start_date = (now - datetime.timedelta(days=30)).strftime("%Y-%m-%d")
    fn_mapper = {
        "unrate": get_unrate_history_data,
        "fedfunds": get_fedfunds_history_data,
    }

    fn = fn_mapper.get(table_name)
    if not fn:
        logging.error("Table %s is not supported.", table_name)
        return

    with financial_engine.connect() as conn:
        for i, row in fn(start_date, end_date, all_data=all_data).items():
            date = str(i).split(" ")[0]
            rate = row

            # check if the data already exists
            if not all_data:
                result = conn.execute(sa.text(
                    f"SELECT * FROM {table_name} WHERE date = :date"
                ), {"date": date})
                if result.fetchone():
                    logging.info("Data for %s of %s already exists.", date, table_name)
                    continue

            data_to_insert.append({
                "date": date,
                "rate": rate
            })

        if data_to_insert:
            conn.execute(sa.text(
                f"INSERT INTO {table_name} (date, rate) VALUES (:date, :rate)",
            ), data_to_insert)

        conn.commit()

        logging.info("Inserted %d rows of data into %s.", len(data_to_insert), table_name)


if __name__ == "__main__":
    # logging.basicConfig(level=logging.INFO)

    # start_date = '2021-01-01'
    # end_date = '2021-01-10'
    # print(get_gold_history_data(start_date, end_date, all_data=True))
    # print(get_oil_history_data(start_date, end_date, all_data=True))
    # print(get_fedfunds_history_data(start_date, end_date, all_data=True))
    # print(get_unrate_history_data(start_date, end_date, all_data=True))
    # print(get_nasdaq_history_data(start_date, end_date, all_data=True))
    refresh_two_cols_history_data("unrate", all_data=True)
    refresh_two_cols_history_data("fedfunds", all_data=True)
    refresh_six_cols_history_data("csi300", all_data=True)
    refresh_six_cols_history_data("gold", all_data=True)
    refresh_six_cols_history_data("crude_oil", all_data=True)
    refresh_six_cols_history_data("nasdaq", all_data=True)
    refresh_six_cols_history_data("dji", all_data=True)