import os
import pickle
import sys
import time

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


BASE_DIR = os.path.dirname(os.path.abspath(__file__))
PROJECT_ROOT = os.path.dirname(BASE_DIR)
sys.path.append(PROJECT_ROOT)

from finetune.config import Config


def main():
    total_start = time.time()
    config = Config()

    output_dir = config.dataset_path
    train_start = pd.Timestamp(config.ddb_train_start_date)
    train_end = pd.Timestamp(config.ddb_train_end_date)
    code_count = config.ddb_code_count
    fetch_code_batch_size = config.ddb_fetch_code_batch_size
    flush_every_chunks = config.ddb_flush_every_chunks

    if fetch_code_batch_size <= 0:
        raise ValueError("config.ddb_fetch_code_batch_size must be greater than 0")
    if flush_every_chunks <= 0:
        raise ValueError("config.ddb_flush_every_chunks must be greater than 0")

    os.makedirs(output_dir, exist_ok=True)
    log_path = os.path.join(output_dir, time.strftime("%Y%m%d_%H%M%S") + ".log")

    def log_print(message):
        print(message)
        with open(log_path, "a", encoding="utf-8") as f:
            f.write(message + "\n")

    session = ddb.session()
    session.connect(config.ddb_host, config.ddb_port, config.ddb_user, config.ddb_password)

    code_list_script = f"""
        exec distinct code
        from {config.ddb_table}
        where trade_date = 2026.04.20
    """
    code_list = session.run(code_list_script)
    if code_list is None or len(code_list) == 0:
        raise RuntimeError("没有取到任何 code")
    code_list = sorted(list(set(code_list.tolist())))[:code_count]
    log_print(f"code count={len(code_list)}")

    query_costs = []
    preprocess_costs = []
    save_costs = []
    part_data = {}
    save_part_idx = 0
    total_chunks = (len(code_list) + fetch_code_batch_size - 1) // fetch_code_batch_size

    for chunk_idx, start_idx in enumerate(range(0, len(code_list), fetch_code_batch_size), start=1):
        code_chunk = code_list[start_idx : start_idx + fetch_code_batch_size]
        code_str = "[" + ",".join([f'"{code}"' for code in code_chunk]) + "]"
        log_print(f"chunk={chunk_idx}/{total_chunks} query {len(code_chunk)} codes")
        query_script = f"""
            select code, trade_date, trade_time, open, high, low, close, vol, amount
            from {config.ddb_table}
            where code in {code_str}
              and trade_date >= {train_start.strftime('%Y.%m.%d')}
              and trade_date < {train_end.strftime('%Y.%m.%d')}
        """

        query_start = time.time()
        chunk_df = session.run(query_script)
        query_cost = time.time() - query_start
        query_costs.append(query_cost)
        log_print(f"query_cost={query_cost:.2f}s")

        chunk_data = {}
        preprocess_cost = 0.0
        if chunk_df is not None and len(chunk_df) > 0:
            preprocess_start = time.time()
            chunk_df = chunk_df.rename(columns={"vol": "volume"})
            chunk_df["timestamps"] = chunk_df["trade_date"] + (
                chunk_df["trade_time"] - chunk_df["trade_time"].dt.normalize()
            )
            chunk_df = chunk_df.drop(columns=["trade_date", "trade_time"], errors="ignore")
            chunk_df = chunk_df.sort_values(["code", "timestamps"]).reset_index(drop=True)
            for code, code_df in chunk_df.groupby("code", sort=False):
                chunk_data[str(code)] = code_df
            preprocess_cost = time.time() - preprocess_start
        preprocess_costs.append(preprocess_cost)
        log_print(f"preprocess_cost={preprocess_cost:.2f}s")

        if len(chunk_data) > 0:
            part_data.update(chunk_data)

        should_flush = (chunk_idx % flush_every_chunks == 0) or (chunk_idx == total_chunks)
        if should_flush:
            save_part_idx += 1
            save_path = os.path.join(output_dir, f"train_part_{save_part_idx:04d}.pkl")
            save_start = time.time()
            with open(save_path, "wb") as f:
                pickle.dump(part_data, f, protocol=pickle.HIGHEST_PROTOCOL)
            save_cost = time.time() - save_start
            save_costs.append(save_cost)
            log_print(
                f"save_part={save_part_idx} chunk_end={chunk_idx} "
                f"save_pkl_cost={save_cost:.2f}s file={os.path.basename(save_path)}"
            )

            part_data = {}

    if len(query_costs) > 0:
        log_print(
            f"query total={sum(query_costs):.2f}s avg={np.mean(query_costs):.2f}s "
            f"min={np.min(query_costs):.2f}s max={np.max(query_costs):.2f}s"
        )
    if len(preprocess_costs) > 0:
        log_print(
            f"preprocess total={sum(preprocess_costs):.2f}s avg={np.mean(preprocess_costs):.2f}s "
            f"min={np.min(preprocess_costs):.2f}s max={np.max(preprocess_costs):.2f}s"
        )
    if len(save_costs) > 0:
        log_print(
            f"save total={sum(save_costs):.2f}s avg={np.mean(save_costs):.2f}s "
            f"min={np.min(save_costs):.2f}s max={np.max(save_costs):.2f}s"
        )
    log_print(f"total cost={time.time() - total_start:.2f}s")


if __name__ == "__main__":
    main()