#!/usr/bin/env python3
"""Tests for capital restore source selection on startup."""

import json
import os
from datetime import datetime, timezone
from unittest.mock import patch

import pytest

from prod.nautilus_event_trader import DolphinLiveTrader


class _MapStub:
    def __init__(self, payloads):
        self._payloads = payloads

    def blocking(self):
        return self

    def get(self, key):
        return self._payloads.get(key)


def _build_trader() -> DolphinLiveTrader:
    trader = DolphinLiveTrader()
    trader._build_engine()
    trader.eng.begin_day(datetime.now(timezone.utc).strftime("%Y-%m-%d"), posture="APEX")
    return trader


def test_restore_prefers_fresher_engine_snapshot_over_stale_latest_nautilus():
    trader = _build_trader()
    trader.eng.capital = 25_000.0

    trader.state_map = _MapStub(
        {
            "latest_nautilus": json.dumps(
                {
                    "capital": 31_049.44,
                    "updated_at": "2026-05-13T10:52:40+00:00",
                }
            ),
            "engine_snapshot": json.dumps(
                {
                    "capital": 33_150.07,
                    "timestamp": "2026-05-13T16:20:38+00:00",
                }
            ),
        }
    )
    trader.pnl_map = _MapStub({})

    with patch.dict(os.environ, {"DOLPHIN_CAPITAL_SEED_STALE_LAG_SEC": "180"}, clear=False):
        trader._restore_capital()

    assert trader.eng.capital == pytest.approx(33_150.07, abs=0.01)
    assert trader._restore_source == "HZ engine_snapshot"


def test_restore_can_force_latest_nautilus_override():
    trader = _build_trader()
    trader.eng.capital = 25_000.0

    trader.state_map = _MapStub(
        {
            "latest_nautilus": json.dumps(
                {
                    "capital": 31_049.44,
                    "updated_at": "2026-05-13T10:52:40+00:00",
                }
            ),
            "engine_snapshot": json.dumps(
                {
                    "capital": 33_150.07,
                    "timestamp": "2026-05-13T16:20:38+00:00",
                }
            ),
        }
    )
    trader.pnl_map = _MapStub({})

    with patch.dict(os.environ, {"DOLPHIN_FORCE_LATEST_NAUTILUS_RESTORE": "1"}, clear=False):
        trader._restore_capital()

    assert trader.eng.capital == pytest.approx(31_049.44, abs=0.01)
    assert trader._restore_source == "HZ latest_nautilus"