from __future__ import annotations

from dataclasses import replace

from .focus_types import AnchorState, FocusCandidate


class AnchorMap:
    def __init__(self, capacity: int = 8) -> None:
        self.capacity = max(1, int(capacity))
        self._anchors: dict[tuple[int, int], AnchorState] = {}

    def clear(self) -> None:
        self._anchors.clear()

    def count(self) -> int:
        return len(self._anchors)

    def states(self) -> list[AnchorState]:
        return list(self._anchors.values())

    def nearest_anchor(self, point_value: tuple[int, int]) -> tuple[AnchorState | None, float]:
        if not self._anchors:
            return None, float("inf")
        best_anchor = None
        best_distance = float("inf")
        for anchor_state in self._anchors.values():
            if anchor_state.point is None:
                continue
            distance_value = _manhattan_distance(anchor_state.point, point_value)
            if distance_value < best_distance or (
                distance_value == best_distance and anchor_state.point < (best_anchor.point if best_anchor is not None and best_anchor.point is not None else (10**9, 10**9))
            ):
                best_anchor = anchor_state
                best_distance = distance_value
        return best_anchor, best_distance

    def touch_candidate(
        self,
        candidate: FocusCandidate,
        tick_number: int,
        policy_name: str,
    ) -> tuple[str, AnchorState]:
        if candidate.point is None:
            anchor_state = AnchorState(
                kind="scalar",
                scalar_value=candidate.scalar_value,
                created_tick=int(tick_number),
                updated_tick=int(tick_number),
                visit_count=1,
                energy=float(candidate.energy),
                policy_name=policy_name,
            )
            return "update", anchor_state

        existing_anchor = self._anchors.get(candidate.point)
        if existing_anchor is None:
            anchor_state = AnchorState(
                kind="spatial",
                point=candidate.point,
                created_tick=int(tick_number),
                updated_tick=int(tick_number),
                visit_count=1,
                energy=float(candidate.energy),
                policy_name=policy_name,
            )
            self._insert_anchor(anchor_state)
            return "create", anchor_state

        anchor_state = replace(
            existing_anchor,
            updated_tick=int(tick_number),
            visit_count=int(existing_anchor.visit_count) + 1,
            energy=max(float(existing_anchor.energy), float(candidate.energy)),
            policy_name=policy_name,
        )
        self._anchors[candidate.point] = anchor_state
        if anchor_state.visit_count > 1:
            return "revisit", anchor_state
        return "update", anchor_state

    def expire_before(self, oldest_allowed_tick: int) -> None:
        expired_keys = [
            point_value
            for point_value, anchor_state in self._anchors.items()
            if int(anchor_state.updated_tick) < int(oldest_allowed_tick)
        ]
        for point_value in expired_keys:
            self._anchors.pop(point_value, None)

    def _insert_anchor(self, anchor_state: AnchorState) -> None:
        if anchor_state.point is None:
            return
        if len(self._anchors) >= self.capacity:
            oldest_point = min(
                self._anchors.keys(),
                key=lambda point_value: (
                    int(self._anchors[point_value].updated_tick),
                    int(point_value[0]),
                    int(point_value[1]),
                ),
            )
            self._anchors.pop(oldest_point, None)
        self._anchors[anchor_state.point] = anchor_state


def _manhattan_distance(left_point: tuple[int, int], right_point: tuple[int, int]) -> float:
    return float(abs(int(left_point[0]) - int(right_point[0])) + abs(int(left_point[1]) - int(right_point[1])))
