"""Backtest simulation orchestration.

Provides BacktestRunner for automating backtest execution with hooks,
progress reporting, and error handling.
"""

from collections import defaultdict
from collections.abc import Callable
from contextlib import suppress
from dataclasses import dataclass, field
from functools import cached_property
from time import monotonic, time
from typing import Any

import pandas as pd
from tqdm import tqdm

from .core import Portfolio, PortfolioState
from .logging import LogEntry, OutputMode, Severity
from .strategy import BaseStrategy


@dataclass
class BacktestResult:
    """Container for backtest results.

    This is a lean data container that holds references to the portfolio and
    strategy after backtest completion, along with execution metadata.

    All analysis (returns, Sharpe ratio, drawdowns, etc.) should be calculated
    by the user from portfolio.history and portfolio.transactions.

    Attributes:
        portfolio: Portfolio instance after backtest completion
        strategy: Strategy instance after backtest completion
        start_period: First period of the backtest
        end_period: Last period of the backtest
        total_periods: Number of periods simulated
        execution_time: Wall-clock time in seconds
        log: Structured log of all events during the backtest
    """

    portfolio: Portfolio
    strategy: BaseStrategy
    start_period: pd.Period
    end_period: pd.Period
    total_periods: int
    execution_time: float
    log: list[LogEntry] = field(default_factory=list)

    def __repr__(self) -> str:
        error_count = len(self.errors)
        error_str = f", {error_count} errors" if error_count else ""
        return (
            f"BacktestResult({self.start_period} to {self.end_period}, "
            f"{self.total_periods} periods, {self.execution_time:.2f}s{error_str})"
        )

    @cached_property
    def log_df(self) -> pd.DataFrame:
        """Lazy DataFrame view of the log."""
        if not self.log:
            return pd.DataFrame(
                columns=[
                    "severity",
                    "timestamp",
                    "period",
                    "source",
                    "message",
                    "data",
                ]
            )
        return pd.DataFrame(
            [
                {
                    "severity": e.severity.name,
                    "timestamp": e.timestamp,
                    "period": e.period,
                    "source": e.source,
                    "message": e.message,
                    "data": e.data,
                }
                for e in self.log
            ]
        )

    @property
    def errors(self) -> list[LogEntry]:
        """Log entries with severity >= ERROR."""
        return [e for e in self.log if e.severity >= Severity.ERROR]

    @property
    def warnings(self) -> list[LogEntry]:
        """Log entries with severity == WARNING."""
        return [e for e in self.log if e.severity == Severity.WARNING]

    @property
    def success(self) -> bool:
        """Returns True if backtest completed without errors."""
        return not any(e.severity >= Severity.ERROR for e in self.log)


class BacktestRunner:
    """Orchestrates backtest execution with hooks for customization.

    This class automates the period-by-period execution loop for backtesting,
    handling income collection, strategy execution, history updates, and
    period advancement.

    The runner enforces the Portfolio state machine:
        COLLECT_INCOME -> TRANSACT -> DONE -> advance_period -> COLLECT_INCOME

    Hooks can be registered for the following events:
        - 'period_start': Called at the beginning of each period
        - 'period_end': Called after update_history() but before advance_period()
        - 'backtest_start': Called before the first period
        - 'backtest_end': Called after the last period
        - 'error': Called when an exception occurs during a period

    Hook callbacks should accept the BacktestRunner instance as their only argument.

    Attributes:
        portfolio: Portfolio being simulated
        strategy: Strategy executing trades
        start_period: First period to simulate (None = portfolio's first period)
        end_period: Last period to simulate (None = portfolio's last period)
        current_period: Current period being processed (None before run)
        strict: If True, any period error re-raises immediately. If False,
            errors are recorded and the backtest continues (lenient mode).

    Example:
        Basic usage::

            runner = BacktestRunner(portfolio, strategy)
            result = runner.run()

        With custom logging::

            def log_cash(runner):
                print(f"{runner.current_period}: {runner.portfolio.cash:,.2f}")

            runner = BacktestRunner(portfolio, strategy)
            runner.register_hook('period_end', log_cash)
            result = runner.run()

        Step-by-step execution::

            runner = BacktestRunner(portfolio, strategy)
            runner.run_period()  # Execute first period
            runner.run_period()  # Execute second period
            # ... etc
    """

    def __init__(
        self,
        portfolio: Portfolio,
        strategy: BaseStrategy,
        *,
        start_period: pd.Period | None = None,
        end_period: pd.Period | None = None,
        strict: bool = True,
    ):
        """Initialize the BacktestRunner.

        Args:
            portfolio: Portfolio instance to simulate
            strategy: Strategy instance to execute trades
            start_period: Optional start period (default: portfolio's current_period)
            end_period: Optional end period (default: last period in universe)
            strict: If True (default), any exception during a period immediately
                re-raises and stops the backtest. If False, errors are recorded
                and the runner attempts to advance to the next period while still
                calling update_history() so portfolio state remains consistent.

        Raises:
            ValueError: If start_period > end_period
            ValueError: If periods are outside the portfolio's universe range
        """
        self.portfolio = portfolio
        self.strategy = strategy

        # Determine period range
        universe_periods = portfolio.asset_universe.get_period_index_range()
        # Resolve start_period: explicit runner arg > strategy.start_period > current
        if start_period is not None:
            resolved_start = start_period
        elif strategy.start_period is not None:
            resolved_start = strategy.start_period
        else:
            resolved_start = portfolio.current_period
        self.start_period = resolved_start
        self.end_period = end_period or universe_periods[-1]

        # Validate period range
        if self.start_period > self.end_period:
            raise ValueError(
                f"start_period {self.start_period} is after "
                f"end_period {self.end_period}"
            )
        if self.start_period not in universe_periods:
            raise ValueError(
                f"start_period {self.start_period} not in asset universe period range"
            )
        if self.end_period not in universe_periods:
            raise ValueError(
                f"end_period {self.end_period} not in asset universe period range"
            )

        # Advance portfolio to start_period if needed
        while portfolio.current_period < self.start_period:
            try:
                # Skip periods before start by just advancing
                if portfolio._states[portfolio.current_period] != PortfolioState.DONE:  # type: ignore[call-overload]
                    # Need to complete the current period first
                    portfolio.collect_income()
                    portfolio.update_history()
                portfolio.advance_period()
            except StopIteration as e:
                raise ValueError(
                    f"Cannot advance to start_period {self.start_period}"
                ) from e

        # Pending initial cash injection — consumed on the first run_period() call
        self._pending_initial_cash: float | None = strategy.initial_cash

        # State tracking
        self.current_period: pd.Period | None = None
        self._log: list[LogEntry] = []
        self._hooks: dict[str, list[Callable]] = defaultdict(list)
        self._periods_completed = 0
        self._strict = strict

    def __repr__(self) -> str:
        strategy_name = type(self.strategy).__name__
        universe_periods = self.portfolio.asset_universe.get_period_index_range()
        start_idx = int(universe_periods.get_loc(self.start_period))  # type: ignore[arg-type]
        end_idx = int(universe_periods.get_loc(self.end_period))  # type: ignore[arg-type]
        total_periods = end_idx - start_idx + 1
        if self._periods_completed == 0 and self.current_period is None:
            status = "not started"
        elif self._periods_completed >= total_periods:
            status = "done"
        else:
            status = "running"
        return (
            f"BacktestRunner(period {self._periods_completed}/{total_periods} | "
            f"strategy={strategy_name} | {status})"
        )

    def _emit(
        self,
        severity: Severity,
        message: str,
        data: dict[str, Any] | None = None,
    ) -> None:
        """Append a structured log entry from the runner."""
        self._log.append(
            LogEntry(
                severity=severity,
                timestamp=monotonic(),
                period=self.portfolio.current_period,
                source="runner",
                message=message,
                data=data,
            )
        )

    def register_hook(self, event: str, callback: Callable) -> None:
        """Register a callback for a specific event.

        Args:
            event: Event name ('period_start', 'period_end', 'backtest_start',
                   'backtest_end', 'error')
            callback: Callable that accepts the BacktestRunner as its only argument

        Example:
            def my_hook(runner):
                print(f"Period: {runner.current_period}")

            runner.register_hook('period_end', my_hook)
        """
        valid_events = {
            "period_start",
            "period_end",
            "backtest_start",
            "backtest_end",
            "error",
        }
        if event not in valid_events:
            raise ValueError(f"Invalid event '{event}'. Must be one of {valid_events}")
        self._hooks[event].append(callback)

    def _trigger_hooks(self, event: str) -> None:
        """Trigger all hooks registered for an event."""
        for callback in self._hooks[event]:
            try:
                callback(self)
            except Exception as exc:
                self._log.append(
                    LogEntry(
                        severity=Severity.WARNING,
                        timestamp=monotonic(),
                        period=self.portfolio.current_period,
                        source="hook",
                        message=(
                            f"Hook '{event}' raised {exc.__class__.__name__}: {exc}"
                        ),
                        data={"event": event, "exception": exc},
                    )
                )

    def run_period(self) -> None:
        """Execute a single period of the backtest.

        This method performs the complete period cycle:
        1. Trigger 'period_start' hooks
        2. Collect income
        3. Execute strategy
        4. Update history
        5. Trigger 'period_end' hooks
        6. Advance to next period

        If an error occurs, it's caught, stored, and 'error' hooks are triggered.
        The portfolio state may be inconsistent after an error.

        Raises:
            StopIteration: If end of backtest period range is reached
            RuntimeError: If portfolio is in an invalid state
        """
        self.current_period = self.portfolio.current_period

        # Check if we've reached the end
        if self.current_period > self.end_period:
            raise StopIteration("End of backtest period range reached")

        self._trigger_hooks("period_start")

        try:
            # Execute the standard period cycle
            self.portfolio.collect_income()

            # Inject initial cash exactly once, on the first active period.
            # The flag is cleared before calling move_cash() to guard against
            # double-injection when strict=False. In that mode (lenient),
            # exceptions during a period are caught and execution continues to
            # the next period. If the flag were still set when move_cash()
            # raised, the next period would attempt a second injection.
            if self._pending_initial_cash is not None:
                cash_to_inject = self._pending_initial_cash
                self._pending_initial_cash = None
                self.portfolio.move_cash(cash_to_inject)

            self.strategy.step()

            # Drain strategy log entries into the runner's log
            self._log.extend(self.strategy._log)
            self.strategy._log.clear()

            self.portfolio.update_history()

            self._trigger_hooks("period_end")

            self._periods_completed += 1

            # Advance to next period
            # If we just completed the end_period, we're done
            if self.current_period >= self.end_period:
                raise StopIteration("Completed all periods")

            # End of portfolio history is expected - suppress StopIteration
            with suppress(StopIteration):
                self.portfolio.advance_period()

        except StopIteration:
            # Let StopIteration pass through - it's how we exit the loop
            raise
        except Exception as e:
            # Store error in structured log
            self._emit(
                Severity.ERROR,
                f"Exception during period {self.current_period}: {e}",
                data={"exception": e},
            )
            self._trigger_hooks("error")

            if self._strict:
                raise

            # Lenient mode: log warning and attempt graceful state recovery.
            self._emit(
                Severity.WARNING,
                f"Lenient mode: skipped error in period {self.current_period}",
            )

            try:
                current_state = self.portfolio._states[self.current_period]  # type: ignore[call-overload]

                # If collect_income didn't finish, force to TRANSACT so that
                # update_history() can still run and record actual portfolio state.
                if current_state == PortfolioState.COLLECT_INCOME:
                    self.portfolio._states[self.current_period] = (  # type: ignore[call-overload]
                        PortfolioState.TRANSACT
                    )

                # update_history() transitions state to DONE and writes the real
                # portfolio snapshot — this avoids forcing DONE without recording.
                period_state = self.portfolio._states[self.current_period]  # type: ignore[call-overload]
                if period_state == PortfolioState.TRANSACT:
                    self.portfolio.update_history()

                self.portfolio.advance_period()
            except StopIteration:
                raise  # End of backtest — normal exit
            except RuntimeError:
                raise RuntimeError(
                    f"Cannot recover from error in period {self.current_period}. "
                    f"Portfolio is in an inconsistent state."
                ) from e

    def run(self, *, output: OutputMode = OutputMode.SILENT) -> BacktestResult:
        """Execute the complete backtest simulation.

        Args:
            output: Controls terminal output. SILENT captures log only,
                SUMMARY prints a one-line summary at end, PROGRESS shows a
                tqdm progress bar plus summary.

        Returns:
            BacktestResult containing portfolio, strategy, and metadata

        Example:
            result = runner.run(output=OutputMode.PROGRESS)
            print(f"Completed {result.total_periods} periods in "
                  f"{result.execution_time:.2f}s")
        """
        # Calculate total periods to run
        universe_periods = self.portfolio.asset_universe.get_period_index_range()
        start_idx: int = universe_periods.get_loc(self.start_period)  # type: ignore[assignment]
        end_idx: int = universe_periods.get_loc(self.end_period)  # type: ignore[assignment]
        total_periods = end_idx - start_idx + 1

        # Trigger start hooks
        self._trigger_hooks("backtest_start")

        # Setup progress bar if requested
        if output == OutputMode.PROGRESS:
            pbar = tqdm(
                total=total_periods,
                desc="Running backtest",
                unit="period",
            )
        else:
            pbar = None

        # Execute the backtest
        start_time = time()
        try:
            while True:
                try:
                    self.run_period()
                    if pbar:
                        pbar.update(1)
                except StopIteration:
                    # Normal completion
                    break
        finally:
            if pbar:
                pbar.close()

        execution_time = time() - start_time

        # Trigger end hooks
        self._trigger_hooks("backtest_end")

        # Check for trade failures and emit summary warning
        if hasattr(self.strategy, "trades_df") and len(self.strategy.trades_df) > 0:
            failed = self.strategy.trades_df[
                self.strategy.trades_df["success"] == False  # noqa: E712
            ]
            if len(failed) > 0:
                total_trades = len(self.strategy.trades_df)
                self._emit(
                    Severity.WARNING,
                    f"Backtest completed with {len(failed)}/{total_trades} "
                    f"failed trades in trades_df",
                    data={"failed_count": len(failed), "total_count": total_trades},
                )

        # Create result
        result = BacktestResult(
            portfolio=self.portfolio,
            strategy=self.strategy,
            start_period=self.start_period,
            end_period=self.end_period,
            total_periods=self._periods_completed,
            execution_time=execution_time,
            log=self._log.copy(),
        )

        # Print summary if requested
        if output in (OutputMode.SUMMARY, OutputMode.PROGRESS):
            error_count = len(result.errors)
            warning_count = len(result.warnings)
            print(
                f"Backtest: {result.total_periods} periods, "
                f"{execution_time:.2f}s, "
                f"{error_count} errors, {warning_count} warnings"
            )

        return result
