# -*- coding: utf-8 -*-
"""网格回撤止盈状态机。

该模块只负责记录每个持仓的最高盈利网格，并判断当前盈亏率是否从
峰值网格回撤。它不包含下单逻辑，由主策略和 Upmax 根据返回的状态决定是否卖出。
"""

from dataclasses import dataclass
from enum import Enum
import math
import threading


class GridState(str, Enum):
    """单次盈亏率观察后的网格状态。"""

    ARMED = "armed"       # 首次记录该持仓的峰值网格
    RAISED = "raised"     # 盈利继续上升，峰值网格已抬高
    RETREAT = "retreat"   # 从峰值网格回撤，应由调用方执行止盈
    STEADY = "steady"     # 仍处于当前峰值网格，继续持有


@dataclass(frozen=True)
class GridObservation:
    """一次网格观察的不可变结果。"""

    state: GridState
    current_grid: int  # 当前盈亏率所处的网格
    peak_grid: int     # 该持仓自观察以来的最高网格


class GridTrailingTracker:
    """按持仓键隔离、线程安全的峰值网格跟踪器。"""

    def __init__(self, step: float = 1.0):
        """
        Args:
            step: 单个网格的盈亏率跨度（百分点），必须大于 0。
        """
        if step <= 0:
            raise ValueError("grid step must be positive")
        self._step = step
        # key 由调用方组成“账户 + 股票代码”，防止多账户状态串扰。
        self._peaks: dict[str, int] = {}
        # 主策略和回调线程可能并发访问，所有峰值读写均在同一把锁内。
        self._lock = threading.Lock()

    def observe(self, position_key: str, pnl_rate: float) -> GridObservation:
        """记录当前盈亏率，并返回相对于历史峰值的状态。"""
        # floor 保证负数盈亏率也按完整网格向下归档。
        current_grid = math.floor(pnl_rate / self._step)
        with self._lock:
            peak_grid = self._peaks.get(position_key)

            # 第一次看到该持仓：建立基准，不触发止盈。
            if peak_grid is None:
                self._peaks[position_key] = current_grid
                return GridObservation(GridState.ARMED, current_grid, current_grid)

            # 进入更高网格：更新峰值，继续持有。
            if current_grid > peak_grid:
                self._peaks[position_key] = current_grid
                return GridObservation(GridState.RAISED, current_grid, current_grid)

            # 跌破峰值网格：报告回撤，但保留峰值直到卖出成功后 clear。
            if current_grid < peak_grid:
                return GridObservation(GridState.RETREAT, current_grid, peak_grid)

            return GridObservation(GridState.STEADY, current_grid, peak_grid)

    def clear(self, position_key: str) -> None:
        """持仓卖出成功后删除峰值，使下次建仓从新状态开始。"""
        with self._lock:
            self._peaks.pop(position_key, None)

    def retain(self, position_keys) -> None:
        """删除已不在券商持仓中的峰值，避免同代码重新开仓继承旧状态。"""
        active = set(position_keys)
        with self._lock:
            self._peaks = {key: value for key, value in self._peaks.items() if key in active}