import argparse
import os
import pickle
import sys
from datetime import datetime
import time

import dolphindb as ddb
import numpy as np
import pandas as pd


PROJECT_ROOT = os.path.dirname(os.path.abspath(__file__))

sys.path.append(PROJECT_ROOT)


DEFAULT_OUTPUT_DIR = os.path.join(PROJECT_ROOT, "2026_bk_data")
START_DATE = pd.Timestamp("2026-01-01 00:00:00")
END_DATE = pd.Timestamp("2026-04-21 00:00:00")

DDB_HOST = "localhost"
DDB_PORT = 8848
DDB_USER = "admin"
DDB_PASSWORD = "123456"
FACTOR_TABLE = 'loadTable("dfs://level2", "factor_table")'
PRICE_TABLE = 'loadTable("dfs://level2", "kline_table")'
INDUSTRY_TABLE = 'loadTable("dfs://level2", "industry_table")'

INDUSTRY_CODE = "801044.SI"


def format_ddb_string_list(items):
    return "[" + ",".join([f'"{item}"' for item in items]) + "]"


def create_ddb_session():
    session = ddb.session()
    session.connect(DDB_HOST, DDB_PORT, DDB_USER, DDB_PASSWORD)
    return session

def fetch_code_list_from_ddb(session):
    script = f"""
        exec distinct code
        from {INDUSTRY_TABLE}
        where l2_code == "{INDUSTRY_CODE}"
    """
    code_df = session.run(script)
    if code_df is None or len(code_df) == 0:
        return []
    return code_df.tolist()

def fetch_industry_data_from_ddb(session, code_list, start_time, end_time):
    if not code_list:
        return {}

    code_filter = format_ddb_string_list(code_list)
    script = f"""
        factor = select factorvalue
        from {FACTOR_TABLE}
        where code in {code_filter}
          and trade_date >= {start_time.strftime('%Y.%m.%d')}
          and trade_date < {end_time.strftime('%Y.%m.%d')}
          and trade_time >= 09:30:00
          and trade_time <= 15:00:00
        pivot by code, trade_date, trade_time, factorname

        price = select code, trade_date, trade_time, close
                from {PRICE_TABLE}
                where code in {code_filter}
                  and trade_date >= {start_time.strftime('%Y.%m.%d')}
                  and trade_date < {end_time.strftime('%Y.%m.%d')}
                  and trade_time >= 09:30:00
                  and trade_time <= 15:00:00

        lj(price, factor, `code`trade_date`trade_time)
    """

    chunk_df = session.run(script)
    if chunk_df is None or len(chunk_df) == 0:
        return {}
    
    chunk_df["trade_time"] = chunk_df["trade_time"].dt.strftime("%H:%M:%S")
    chunk_df["timestamps"] = pd.to_datetime(
        chunk_df["trade_date"].astype(str) + " " + chunk_df["trade_time"].astype(str)
    )
    chunk_df = chunk_df.drop(columns=["trade_date", "trade_time"], errors="ignore")
    chunk_df = chunk_df.sort_values(["code", "timestamps"]).reset_index(drop=True)

    grouped = {}
    for code, code_df in chunk_df.groupby("code", sort=False):
        grouped[str(code)] = code_df.drop(columns=["code"], errors="ignore").reset_index(drop=True)
    return grouped


def save_pkl(obj, path):
    os.makedirs(os.path.dirname(path), exist_ok=True)
    with open(path, "wb") as f:
        pickle.dump(obj, f)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--output_dir", type=str, default=DEFAULT_OUTPUT_DIR)

    args = parser.parse_args()

    total_start = time.time()
    out_dir = args.output_dir
    os.makedirs(out_dir, exist_ok=True)

    start_time = START_DATE
    end_time = END_DATE

    session = create_ddb_session()

    code_list = fetch_code_list_from_ddb(session)
    print(f"code count: {len(code_list)}")
    if len(code_list) == 0:
        raise RuntimeError("没有取到任何 code")

    full_data = fetch_industry_data_from_ddb(session, code_list, start_time, end_time)
    if not full_data:
        raise RuntimeError("没有取到任何行业数据")

    sample_code = next(iter(full_data))
    exclude_columns = {"code", "close", "timestamps"}
    factor_columns = [
        col for col in full_data[sample_code].columns.to_list() if col not in exclude_columns
    ]
    print(f"factor dim: {len(factor_columns)}")
    print(f"factor columns example: {factor_columns[:10]}")

    print(f"fetched grouped code data: {len(full_data)}")

    output_file = os.path.join(out_dir, f"{INDUSTRY_CODE}.pkl")
    save_pkl(full_data, output_file)
    print(f"saved pkl file to {output_file}")

    total_cost = time.time() - total_start
    print(f"total cost={total_cost:.2f}s")
    print("done")


if __name__ == "__main__":
    main()