import logging
import time
from pathlib import Path
from typing import cast

import pandas as pd
from numpy import inf

from signalab.account import Account, SimpleMarketData
from signalab.analytics.report import AccountReport, Report
from signalab.data.market_data import MarketData
from signalab.indicator.base import Indicator
from signalab.strategy.base import Strategy

logger = logging.getLogger(__name__)


class BacktestEngine:
    """Backtesting orchestrator for indicators, strategies, and accounts.

    Notes
    -----
    The engine aligns indicator outputs to tick data, executes strategy logic
    for each row, and records account metrics over time.
    """

    def __init__(
        self,
        verbose: bool = True,
        progress_log_interval: pd.Timedelta | None = None,
        results_path: str | Path | None = None,
    ):
        """Initialize engine runtime defaults.

        Parameters
        ----------
        verbose : bool, default=True
            Whether to emit periodic progress logs.
        progress_log_interval : pd.Timedelta | None, default=pd.Timedelta(hours=6)
            Market-time interval used for progress logging. ``None`` disables
            progress logs.
        results_path : str | Path | None, default=None
            Optional folder path where run result dataframes are saved as CSV.
            If ``None``, results are not written to disk.
        """
        if progress_log_interval is None:
            progress_log_interval = pd.Timedelta(hours=6)
        if progress_log_interval is not None and progress_log_interval <= pd.Timedelta(
            0
        ):
            raise ValueError("progress_log_interval must be positive or None")

        self.verbose = verbose
        self.progress_log_interval = progress_log_interval
        self.results_path = Path(results_path) if results_path is not None else None

    def _save_results(self, results: dict[str, pd.DataFrame]) -> None:
        """Persist run outputs to disk when ``results_path`` is configured."""
        if self.results_path is None:
            return

        self.results_path.mkdir(parents=True, exist_ok=True)
        for name, df in results.items():
            safe_name = name.lower().replace(" ", "_")
            df.to_csv(self.results_path / f"{safe_name}.csv")

    def _setup_accounts(
        self,
        individual_accounts: bool,
        accounts: Account | dict[str, Account] | list[Account] | None,
        strategies: list[Strategy],
    ) -> dict[str, Account]:
        """Normalize account input into a per-strategy account mapping.

        Parameters
        ----------
        individual_accounts : bool
            Whether each strategy should operate on a distinct account.
        accounts : Account | dict[str, Account] | list[Account] | None
            Account input supplied by the caller.
        strategies : list[Strategy]
            Active strategies for this run.

        Returns
        -------
        dict[str, Account]
            Mapping from strategy label to account instance.

        Raises
        ------
        ValueError
            If provided account mapping/list does not match strategies.
        """
        if accounts is None:
            accounts = Account(inf, label="common_account")

        if isinstance(accounts, Account):
            if individual_accounts:
                return {
                    strategy.label: accounts.copy_empty() for strategy in strategies
                }
            return {strategy.label: accounts for strategy in strategies}
        if isinstance(accounts, dict):
            key_labels = {key.lower().replace(" ", "_") for key in accounts.keys()}
            strategy_labels = {strategy.label for strategy in strategies}
            if key_labels != strategy_labels:
                raise ValueError(
                    "Account keys do not match strategy labels. "
                    f"Expected: {strategy_labels}, got: {key_labels}"
                )
            return {
                key.lower().replace(" ", "_"): account
                for key, account in accounts.items()
            }
        if isinstance(accounts, list):
            if len(accounts) != len(strategies):
                raise ValueError(
                    "Number of accounts does not match number of strategies. "
                    f"Expected: {len(strategies)}, got: {len(accounts)}"
                )
            return {
                strategy.label: account
                for strategy, account in zip(strategies, accounts, strict=True)
            }

        raise ValueError(
            "Invalid accounts parameter. Must be an Account instance, a dict of accounts, or a list of accounts."
        )

    def _compute_indicators(
        self, record_df: pd.DataFrame, indicators: list[Indicator], data: MarketData
    ) -> pd.DataFrame:
        """Compute all indicators and merge their outputs into the record frame.

        Parameters
        ----------
        record_df : pd.DataFrame
            Base tick dataframe used by the simulation loop.
        indicators : list[Indicator]
            Indicators to compute and merge.
        data : MarketData
            Source market data object.

        Returns
        -------
        pd.DataFrame
            Enriched dataframe containing bid/ask plus indicator columns.
        """
        # new_rows_df accumulates all extra rows (candle closes not in tick index)
        # keyed by timestamp — signals from multiple indicators at the same
        # timestamp are merged column-by-column into a single row here, so the
        # final concat never produces duplicate index entries.
        new_rows_df: pd.DataFrame | None = None
        signal_updates: dict[str, pd.Series] = {}
        continuous_dfs: list[pd.DataFrame] = []

        for indicator in indicators:
            logger.info("Computing indicator: %s", indicator.name)
            indicator_label = indicator.label or ""
            indicator_df = indicator.compute(data)
            indicator_df.index = pd.DatetimeIndex(indicator_df.index)
            indicator_df.index.name = "timestamp"

            if indicator_label:
                indicator_label = indicator_label.lower().replace(" ", "_")
                indicator_df = indicator_df.rename(
                    columns={
                        col: f"{indicator_label}_{col}" for col in indicator_df.columns
                    }
                )

            indicator_df = indicator_df.sort_index()

            all_cols = list(indicator_df.columns)
            signal_cols = [col for col in all_cols if col.endswith("_signal")]
            continuous_cols = [col for col in all_cols if col not in signal_cols]

            # ── Stash continuous columns for a single merge pass at the end ───
            if continuous_cols:
                continuous_dfs.append(indicator_df[continuous_cols])

            # ── Signal columns: insert/update at exact candle-close timestamp ──
            if signal_cols:
                sig_df = indicator_df[signal_cols].copy()
                sig_df = sig_df[sig_df.any(axis=1)]  # drop all-zero rows

                new_mask = ~sig_df.index.isin(record_df.index)
                existing_mask = ~new_mask

                # Rows that match an existing tick: collect for later .loc update
                if existing_mask.any():
                    for col in signal_cols:
                        s = sig_df.loc[existing_mask, col]
                        signal_updates[col] = (
                            pd.concat([signal_updates[col], s])
                            if col in signal_updates
                            else s
                        )

                # Rows with no matching tick: merge into new_rows_df
                if new_mask.any():
                    new_sig = sig_df.loc[new_mask, signal_cols]

                    # Forward-fill bid/ask from the last known tick
                    locs = record_df.index.searchsorted(new_sig.index, side="right") - 1
                    locs = locs.clip(0, len(record_df.index) - 1)

                    # Build a stub with bid/ask + this indicator's signal columns
                    stub = pd.DataFrame(index=new_sig.index)
                    stub.index.name = "timestamp"
                    stub["bid"] = record_df["bid"].iloc[locs].values
                    stub["ask"] = record_df["ask"].iloc[locs].values
                    for col in signal_cols:
                        stub[col] = new_sig[col].values

                    if new_rows_df is None:
                        new_rows_df = stub
                    else:
                        # Merge stub into new_rows_df:
                        # - timestamps already present → assign new signal columns in place
                        # - timestamps not yet present → append as new rows
                        shared_idx = new_rows_df.index.intersection(stub.index)
                        if not shared_idx.empty:
                            for col in signal_cols:
                                if col in stub.columns:
                                    new_rows_df.loc[shared_idx, col] = stub.loc[
                                        shared_idx, col
                                    ]

                        only_in_stub = stub.index.difference(new_rows_df.index)
                        if not only_in_stub.empty:
                            new_rows_df = pd.concat(
                                [new_rows_df, stub.loc[only_in_stub]]
                            )

        # ── Insert all new rows in one shot ──────────────────────────────────
        if new_rows_df is not None:
            new_rows_df = new_rows_df.sort_index()
            record_df = pd.concat([record_df, new_rows_df]).sort_index()

        # ── Merge all continuous columns over the final, complete index ───────
        for cont_df in continuous_dfs:
            record_df = pd.merge_asof(
                record_df.reset_index(),
                cont_df.reset_index(),
                on="timestamp",
                direction="backward",
            ).set_index("timestamp")

        # ── Apply signal updates to existing tick rows ────────────────────────
        for col, series in signal_updates.items():
            if col not in record_df.columns:
                record_df[col] = 0
            record_df.loc[series.index, col] = series.values

        # ── Fill any remaining signal NaNs with 0 ────────────────────────────
        all_signal_cols = [col for col in record_df.columns if col.endswith("_signal")]
        record_df[all_signal_cols] = record_df[all_signal_cols].fillna(0).astype(int)

        return record_df

    def run(
        self,
        data: MarketData,
        indicators: list[Indicator],
        strategies: list[Strategy],
        from_date: pd.Timestamp | None = None,
        to_date: pd.Timestamp | None = None,
        individual_accounts: bool = False,
        accounts: Account | dict[str, Account] | list[Account] | None = None,
    ) -> Report:
        """Execute a full backtest run.

        Parameters
        ----------
        data : MarketData
            Source market data.
        indicators : list[Indicator]
            Indicator instances used to generate features and signals.
        strategies : list[Strategy]
            Strategies executed on each market row.
        from_date : pd.Timestamp | None, default=None
            Inclusive lower date bound.
        to_date : pd.Timestamp | None, default=None
            Inclusive upper date bound.
        individual_accounts : bool, default=False
            If ``True``, each strategy receives a separate account.
        accounts : Account | dict[str, Account] | list[Account] | None, default=None
            Account configuration input.

        Returns
        -------
        dict[str, pd.DataFrame]
            Result dictionary with ``indicators`` dataframe plus one dataframe
            per account label.
        """
        verbose = self.verbose
        progress_log_interval = self.progress_log_interval

        total_start_ns = time.perf_counter_ns()
        logger.info(
            "Running simulation from %s to %s",
            str(from_date),
            str(to_date),
        )
        data = data.filter(from_date, to_date)

        logger.info("Number of ticks: %d", len(data.ticks))

        tick_index = pd.DatetimeIndex(data.ticks.index)

        record_df = pd.DataFrame(index=tick_index)
        record_df["ask"] = data.ticks["ask"]
        record_df["bid"] = data.ticks["bid"]
        record_df.index.name = "timestamp"

        logger.info(
            "Indicators: \n    %s", "\n    ".join(str(ind) for ind in indicators)
        )

        record_df = self._compute_indicators(record_df, indicators, data)

        logger.info("Setting up accounts")
        accounts = self._setup_accounts(individual_accounts, accounts, strategies)

        for strategy in strategies:
            account = accounts[strategy.label]
            if individual_accounts:
                account.label = strategy.label
            strategy.set_account(accounts[strategy.label])

        account_set = set(accounts.values())
        for account in account_set:
            account.init_record(record_df)

        assert not individual_accounts or len(account_set) == len(strategies), (
            "Individual accounts flag is set but multiple accounts provided."
        )

        logger.info("Preparing market data")
        records_iter = list(record_df.itertuples())
        total_ticks = len(records_iter)

        logger.info("Starting backtest loop")
        loop_start_ns = time.perf_counter_ns()
        next_progress_log_ts: pd.Timestamp | None = None
        if progress_log_interval is not None and total_ticks > 0:
            first_ts = cast(pd.Timestamp, records_iter[0].Index)
            next_progress_log_ts = first_ts + progress_log_interval

        for i, current_market in enumerate(records_iter):
            prev_market = records_iter[i - 1] if i > 0 else None

            day_changed = (
                prev_market is not None
                and cast(pd.Timestamp, prev_market.Index).day
                != cast(pd.Timestamp, current_market.Index).day
            )
            days_delta = 0
            if day_changed:
                assert prev_market is not None
                timedelta = cast(pd.Timestamp, current_market.Index) - cast(
                    pd.Timestamp, prev_market.Index
                )
                days_delta = max(timedelta.days, 1)

            for account in account_set:
                account._set_current_market(cast(SimpleMarketData, current_market))
                account.check_liquidation()
                if day_changed:
                    account.day_change(delta=days_delta)

            for strategy in strategies:
                strategy._set_current_market(current_market)
                strategy._set_all_accounts(account_set)
                strategy.perform(current_market)

            for account in account_set:
                account.update_record(i)

            if not verbose:
                continue

            current_ts = cast(pd.Timestamp, current_market.Index)
            if next_progress_log_ts is not None and current_ts >= next_progress_log_ts:
                now_ns = time.perf_counter_ns()
                processed = i + 1
                elapsed_ns = now_ns - loop_start_ns
                avg_tick_ns = elapsed_ns / processed

                logger.info(
                    "Processed %.2f%% | elapsed=%.3f s, avg_tick=%.0f ns",
                    processed * 100.0 / total_ticks if total_ticks else 100.0,
                    elapsed_ns / 1e9,
                    avg_tick_ns,
                )
                if progress_log_interval is not None:
                    while (
                        next_progress_log_ts is not None
                        and current_ts >= next_progress_log_ts
                    ):
                        next_progress_log_ts += progress_log_interval

        loop_elapsed_ns = time.perf_counter_ns() - loop_start_ns
        total_elapsed_ns = time.perf_counter_ns() - total_start_ns
        avg_tick_ns = loop_elapsed_ns / total_ticks if total_ticks else 0.0

        logger.info(
            "Backtest timing: total=%d ns (%.3f s), loop=%d ns (%.3f s), approx_per_tick=%.0f ns over %d ticks",
            total_elapsed_ns,
            total_elapsed_ns / 1e9,
            loop_elapsed_ns,
            loop_elapsed_ns / 1e9,
            avg_tick_ns,
            total_ticks,
        )
        logger.info("Backtest completed.")

        return Report(
            indicators_df=record_df,
            accounts_reports={
                account.label: AccountReport(account) for account in account_set
            },
        )
