import logging
from typing import cast

import pandas as pd
from numpy import inf

from signalab.account import Account, SimpleMarketData
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:
    def _setup_accounts(
        self,
        individual_accounts: bool,
        accounts: Account | dict[str, Account] | list[Account] | None,
        strategies: list[Strategy],
    ) -> dict[str, Account]:
        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:
        # 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,
    ) -> dict[str, pd.DataFrame]:
        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())

        logger.info("Starting backtest loop")
        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.check_liquidation(cast(SimpleMarketData, current_market))
                if day_changed:
                    account.day_change(delta=days_delta)

            for strategy in strategies:
                strategy.perform(
                    current_market,
                    account_set,
                )

            for account in account_set:
                account.update_record(i, cast(SimpleMarketData, current_market))

            if (
                prev_market is not None
                and cast(pd.Timestamp, current_market.Index).hour
                != cast(pd.Timestamp, prev_market.Index).hour
            ):
                logger.info(
                    "Processed %d records up to %s",
                    i + 1,
                    str(current_market.Index),
                )

        logger.info("Backtest completed.")

        return {
            "indicators": record_df,
            **{account.label: account.get_record_df() for account in account_set},
        }
