import argparse
import hjson
import importlib

from typing import Any, Dict, List

import pandas as pd

from strategies import SlidingFitStrategy, StaticFitStrategy
from tools.data_loader import load_market_data
from tools.trading_pair import TradingPair
from results import BacktestResult


def load_config(config_path: str) -> Dict:
    with open(config_path, "r") as f:
        config = hjson.load(f)
    return config


def run_all_pairs(
    config: Dict, datafile: str, price_column: str, bt_result: BacktestResult, strategy
) -> None:

    def _create_pairs(config: Dict) -> List[TradingPair]:
        nonlocal datafile
        instruments = config["instruments"]
        all_indexes = range(len(instruments))
        unique_index_pairs = [(i, j) for i in all_indexes for j in all_indexes if i < j]
        pairs = []
        market_data_df = load_market_data(
            f'{config["data_directory"]}/{datafile}', config=config
        )
        for a_index, b_index in unique_index_pairs:
            pair = TradingPair(
                market_data=market_data_df,
                symbol_a=instruments[a_index],
                symbol_b=instruments[b_index],
                price_column=price_column,
            )
            pairs.append(pair)
        return pairs

    pairs_trades = []
    for pair in _create_pairs(config):
        single_pair_trades = strategy.run_pair(
            pair=pair, config=config, bt_result=bt_result
        )
        if single_pair_trades is not None and len(single_pair_trades) > 0:
            pairs_trades.append(single_pair_trades)
    # Check if result_list has any data before concatenating
    if len(pairs_trades) == 0:
        print("No trading signals found for any pairs")
        return None

    result = pd.concat(pairs_trades, ignore_index=True)
    result["time"] = pd.to_datetime(result["time"])
    result = result.set_index("time").sort_index()

    bt_result.collect_single_day_results(result)
    # BacktestResults.print_single_day_results()


def main() -> None:
    parser = argparse.ArgumentParser(description="Run pairs trading backtest.")
    parser.add_argument(
        "--config", type=str, required=True, help="Path to the configuration file."
    )
    args = parser.parse_args()

    CONFIG = load_config(args.config)

    # Dynamically instantiate strategy class
    strategy_class_name = CONFIG.get("strategy_class", "strategies.StaticFitStrategy")
    module_name, class_name = strategy_class_name.rsplit(".", 1)
    module = importlib.import_module(module_name)
    STRATEGY = getattr(module, class_name)()

    # Initialize a dictionary to store all trade results
    all_results: Dict[str, Dict[str, Any]] = {}
    bt_results = BacktestResult(config=CONFIG)

    # Process each data file
    price_column = CONFIG["price_column"]
    for datafile in CONFIG["datafiles"]:
        print(f"\n====== Processing {datafile} ======")

        # Clear the TRADES global dictionary and reset unrealized PnL for the new file
        bt_results.clear_trades()

        # Process data for this file
        try:
            run_all_pairs(
                config=CONFIG,
                datafile=datafile,
                price_column=price_column,
                bt_result=bt_results,
                strategy=STRATEGY,
            )

            # Store results with file name as key
            filename = datafile.split("/")[-1]
            all_results[filename] = {"trades": bt_results.trades.copy()}

            print(f"Successfully processed {filename}")

        except Exception as e:
            print(f"Error processing {datafile}: {str(e)}")

    # Calculate and print results
    bt_results.calculate_returns(all_results)
    bt_results.print_grand_totals()
    bt_results.print_outstanding_positions()


if __name__ == "__main__":
    main()