import pandas as pd
from loguru import logger
from chart import resample_data, QuantChart
from db_config import DatabaseManager, DatabaseConfig


def get_kline(code: str) -> list:
    """
    获取所有指数代码
    :return:
    """
    env = "online"
    env = "daily"
    db_config = DatabaseConfig(env=env)
    logger.info(f"数据库连接: {db_config.connection_string}")

    db_manager = DatabaseManager(db_config)
    sql = f"SELECT date as time, open, high, low, close, volume FROM public.index_kline where code='{code}' order by date;"
    res = db_manager.execute_query(sql)
    data_list = [dict(item) for item in res]
    df = pd.DataFrame(data_list)
    df["time"] = pd.to_datetime(df["time"])
    num_cols = ["open", "high", "low", "close", "volume"]
    for col in num_cols:
        if col in df.columns:
            df[col] = pd.to_numeric(df[col], errors="coerce").astype(float)
    return df


if __name__ == "__main__":
    symbol = "399986"
    timeframe = "1D"

    df = pd.read_csv(
        "/Users/aszer/Documents/vscode/etf/data/index_all_stock.csv",
        encoding="utf-8-sig",
    )
    name = df.loc[df["代码"] == symbol, "名称"].values[0]

    df = get_kline(code=symbol)
    df = resample_data(df, timeframe)
    # df['buy'] = df['time'].apply(lambda x: 1 if pd.to_datetime(x).day % 2 == 1 else 0)
    # df['sell'] = df['time'].apply(lambda x: 1 if pd.to_datetime(x).day % 2 == 0 else 0)
    quant_chart = QuantChart()
    quant_chart.plot_chart(df=df, symbol=symbol, name=name, timeframe=timeframe)