"""State management primitives for rate limiting."""

from __future__ import annotations

import json
import threading
from collections import defaultdict
from datetime import datetime, timedelta, timezone
from pathlib import Path
from typing import Any
from weakref import WeakSet

GLOBAL_AGENT_ID = "__global__"
DEFAULT_SESSION_ID = "__default__"


def _to_datetime_utc(timestamp: datetime | str | None) -> datetime:
    if timestamp is None:
        return datetime.now(timezone.utc)
    if isinstance(timestamp, datetime):
        if timestamp.tzinfo is None:
            return timestamp.replace(tzinfo=timezone.utc)
        return timestamp.astimezone(timezone.utc)
    parsed = datetime.fromisoformat(timestamp)
    if parsed.tzinfo is None:
        return parsed.replace(tzinfo=timezone.utc)
    return parsed.astimezone(timezone.utc)


class RateLimitTracker:
    """Thread-safe in-memory counter with optional JSONL persistence."""

    _TRACKERS_BY_PATH: dict[str, WeakSet["RateLimitTracker"]] = {}

    def __init__(self, persist_path: str | Path | None = ".orchesis/state.jsonl"):
        self.persist_path = Path(persist_path) if persist_path is not None else None
        # DESIGN NOTE: Single lock is sufficient for Phase 1/2.
        # Under 50+ concurrent agents, consider per-agent locks
        # using a dict[str, threading.Lock] with lazy initialization.
        # Current benchmarks show <1ms p99 with single lock at 500 concurrent.
        self._lock = threading.Lock()
        self._events: dict[tuple[str, str, str], list[datetime]] = defaultdict(list)
        self._spend_events: dict[tuple[str, str], list[tuple[datetime, float]]] = defaultdict(list)
        self._spend_totals: dict[tuple[str, str], float] = defaultdict(float)
        self._write_buffer: list[str] = []
        self._buffer_limit: int = 50
        if self.persist_path is not None:
            path_key = str(self.persist_path)
            peers = self._TRACKERS_BY_PATH.get(path_key)
            if peers is not None:
                for peer in list(peers):
                    peer.flush()
            self._load_existing()
            self._TRACKERS_BY_PATH.setdefault(path_key, WeakSet()).add(self)

    def _load_existing(self) -> None:
        if self.persist_path is None or not self.persist_path.exists():
            return
        for line in self.persist_path.read_text(encoding="utf-8").splitlines():
            if not line.strip():
                continue
            try:
                payload = json.loads(line)
            except json.JSONDecodeError:
                continue
            event = payload.get("event", "rate")
            tool = payload.get("tool")
            ts = payload.get("timestamp")
            agent = payload.get("agent_id", GLOBAL_AGENT_ID)
            session = payload.get("session_id", DEFAULT_SESSION_ID)
            cost = payload.get("cost", 0.0)
            if event == "spend" and isinstance(ts, str):
                try:
                    dt = _to_datetime_utc(ts)
                    agent_id = agent if isinstance(agent, str) and agent else GLOBAL_AGENT_ID
                    session_id = (
                        session if isinstance(session, str) and session else DEFAULT_SESSION_ID
                    )
                    numeric_cost = float(cost) if isinstance(cost, int | float) else 0.0
                    spend_key = (agent_id, session_id)
                    self._spend_events[spend_key].append((dt, numeric_cost))
                    self._spend_totals[spend_key] += numeric_cost
                except ValueError:
                    continue
                continue
            if isinstance(tool, str) and isinstance(ts, str):
                try:
                    dt = _to_datetime_utc(ts)
                    agent_id = agent if isinstance(agent, str) and agent else GLOBAL_AGENT_ID
                    session_id = (
                        session if isinstance(session, str) and session else DEFAULT_SESSION_ID
                    )
                    self._events[(agent_id, session_id, tool)].append(dt)
                    # Backward compatibility for legacy state entries carrying cost.
                    numeric_cost = float(cost) if isinstance(cost, int | float) else 0.0
                    if numeric_cost != 0.0:
                        spend_key = (agent_id, session_id)
                        self._spend_events[spend_key].append((dt, numeric_cost))
                        self._spend_totals[spend_key] += numeric_cost
                except ValueError:
                    continue

    def _buffer_write(self, payload: dict[str, Any]) -> None:
        if self.persist_path is None:
            return
        self._write_buffer.append(json.dumps(payload, ensure_ascii=False))
        if len(self._write_buffer) >= self._buffer_limit:
            self._flush_buffer()

    def _flush_buffer(self) -> None:
        if not self._write_buffer or self.persist_path is None:
            return
        self.persist_path.parent.mkdir(parents=True, exist_ok=True)
        with self.persist_path.open("a", encoding="utf-8") as file:
            file.write("\n".join(self._write_buffer) + "\n")
        self._write_buffer.clear()

    def _prune(self, key: tuple[str, str, str], now: datetime, window_seconds: int) -> None:
        threshold = now - timedelta(seconds=window_seconds)
        self._events[key] = [x for x in self._events[key] if x >= threshold]

    def _prune_spend(self, key: tuple[str, str], now: datetime, window_seconds: int) -> None:
        threshold = now - timedelta(seconds=window_seconds)
        self._spend_events[key] = [
            (ts, cost) for ts, cost in self._spend_events[key] if ts >= threshold
        ]
        self._spend_totals[key] = sum(cost for _, cost in self._spend_events[key])

    def record(
        self,
        tool_name: str,
        timestamp: datetime | str | None = None,
        agent_id: str = GLOBAL_AGENT_ID,
        session_id: str = DEFAULT_SESSION_ID,
    ) -> None:
        ts = _to_datetime_utc(timestamp)
        safe_agent_id = agent_id if isinstance(agent_id, str) and agent_id else GLOBAL_AGENT_ID
        safe_session_id = (
            session_id if isinstance(session_id, str) and session_id else DEFAULT_SESSION_ID
        )
        with self._lock:
            key = (safe_agent_id, safe_session_id, tool_name)
            self._events[key].append(ts)
            self._buffer_write(
                {
                    "event": "rate",
                    "tool": tool_name,
                    "timestamp": ts.isoformat(),
                    "agent_id": safe_agent_id,
                    "session_id": safe_session_id,
                }
            )

    def record_spend(
        self,
        agent_id: str,
        cost: float,
        timestamp: datetime | str | None = None,
        session_id: str = DEFAULT_SESSION_ID,
    ) -> None:
        ts = _to_datetime_utc(timestamp)
        safe_agent_id = agent_id if isinstance(agent_id, str) and agent_id else GLOBAL_AGENT_ID
        safe_session_id = (
            session_id if isinstance(session_id, str) and session_id else DEFAULT_SESSION_ID
        )
        safe_cost = float(cost) if isinstance(cost, int | float) else 0.0
        with self._lock:
            spend_key = (safe_agent_id, safe_session_id)
            self._spend_events[spend_key].append((ts, safe_cost))
            self._spend_totals[spend_key] += safe_cost
            self._buffer_write(
                {
                    "event": "spend",
                    "agent_id": safe_agent_id,
                    "session_id": safe_session_id,
                    "timestamp": ts.isoformat(),
                    "cost": safe_cost,
                }
            )

    def check_budget_and_record(
        self,
        agent_id: str,
        cost: float,
        daily_budget: float,
        *,
        window_seconds: int = 86400,
        timestamp: datetime | str | None = None,
        session_id: str = DEFAULT_SESSION_ID,
    ) -> bool:
        """Atomically check budget and record spend when within limit."""
        current = _to_datetime_utc(timestamp)
        safe_agent_id = agent_id if isinstance(agent_id, str) and agent_id else GLOBAL_AGENT_ID
        safe_session_id = (
            session_id if isinstance(session_id, str) and session_id else DEFAULT_SESSION_ID
        )
        safe_cost = float(cost) if isinstance(cost, int | float) else 0.0
        safe_budget = float(daily_budget) if isinstance(daily_budget, int | float) else 0.0
        if safe_cost <= 0:
            return False
        with self._lock:
            spend_key = (safe_agent_id, safe_session_id)
            events = self._spend_events[spend_key]
            threshold = current - timedelta(seconds=window_seconds)
            if events and events[0][0] < threshold:
                self._prune_spend(spend_key, current, window_seconds)
            spent = self._spend_totals[spend_key]
            projected = spent + safe_cost
            if projected > safe_budget:
                return True
            self._spend_events[spend_key].append((current, safe_cost))
            self._spend_totals[spend_key] = projected
            self._buffer_write(
                {
                    "event": "spend",
                    "agent_id": safe_agent_id,
                    "session_id": safe_session_id,
                    "timestamp": current.isoformat(),
                    "cost": safe_cost,
                }
            )
            return False

    def get_count(
        self,
        tool_name: str,
        window_seconds: int,
        *,
        agent_id: str = GLOBAL_AGENT_ID,
        session_id: str = DEFAULT_SESSION_ID,
        now: datetime | str | None = None,
    ) -> int:
        current = _to_datetime_utc(now)
        safe_agent_id = agent_id if isinstance(agent_id, str) and agent_id else GLOBAL_AGENT_ID
        safe_session_id = (
            session_id if isinstance(session_id, str) and session_id else DEFAULT_SESSION_ID
        )
        with self._lock:
            key = (safe_agent_id, safe_session_id, tool_name)
            events = self._events[key]
            threshold = current - timedelta(seconds=window_seconds)
            if events and events[0] >= threshold:
                return len(events)
            if len(events) > window_seconds * 2:
                self._prune(key, current, window_seconds)
                return len(self._events[key])
            return sum(1 for x in events if x >= threshold)

    def is_over_limit(
        self,
        tool_name: str,
        max_requests: int,
        window_seconds: int,
        *,
        agent_id: str = GLOBAL_AGENT_ID,
        session_id: str = DEFAULT_SESSION_ID,
        now: datetime | str | None = None,
    ) -> bool:
        return (
            self.get_count(
                tool_name,
                window_seconds,
                now=now,
                agent_id=agent_id,
                session_id=session_id,
            )
            >= max_requests
        )

    def check_and_record(
        self,
        tool_name: str,
        max_requests: int,
        window_seconds: int,
        *,
        agent_id: str = GLOBAL_AGENT_ID,
        session_id: str = DEFAULT_SESSION_ID,
        timestamp: datetime | str | None = None,
    ) -> bool:
        current = _to_datetime_utc(timestamp)
        safe_agent_id = agent_id if isinstance(agent_id, str) and agent_id else GLOBAL_AGENT_ID
        safe_session_id = (
            session_id if isinstance(session_id, str) and session_id else DEFAULT_SESSION_ID
        )
        with self._lock:
            key = (safe_agent_id, safe_session_id, tool_name)
            threshold = current - timedelta(seconds=window_seconds)
            events = self._events[key]
            if events and events[0] >= threshold:
                over_limit = len(events) >= max_requests
            elif len(events) > window_seconds * 2:
                self._prune(key, current, window_seconds)
                over_limit = len(self._events[key]) >= max_requests
            else:
                over_limit = sum(1 for x in events if x >= threshold) >= max_requests
            self._events[key].append(current)
            self._buffer_write(
                {
                    "event": "rate",
                    "tool": tool_name,
                    "timestamp": current.isoformat(),
                    "agent_id": safe_agent_id,
                    "session_id": safe_session_id,
                }
            )
            return over_limit

    def get_agent_budget_spent(
        self,
        agent_id: str,
        window_seconds: int = 86400,
        *,
        session_id: str = DEFAULT_SESSION_ID,
        now: datetime | str | None = None,
    ) -> float:
        current = _to_datetime_utc(now)
        safe_agent_id = agent_id if isinstance(agent_id, str) and agent_id else GLOBAL_AGENT_ID
        safe_session_id = (
            session_id if isinstance(session_id, str) and session_id else DEFAULT_SESSION_ID
        )
        spend_key = (safe_agent_id, safe_session_id)
        with self._lock:
            events = self._spend_events[spend_key]
            if not events:
                return 0.0

            threshold = current - timedelta(seconds=window_seconds)
            if events[0][0] >= threshold:
                return self._spend_totals[spend_key]

            if len(events) > window_seconds * 2:
                self._prune_spend(spend_key, current, window_seconds)
                return self._spend_totals[spend_key]

            return sum(cost for ts, cost in events if ts >= threshold)

    def get_tools(self) -> list[str]:
        with self._lock:
            return sorted({tool for (_, _, tool) in self._events.keys()})

    def get_agents(self) -> list[str]:
        with self._lock:
            return sorted({agent for (agent, _, _) in self._events.keys()})

    def flush(self) -> None:
        """Force buffered state writes to disk."""
        with self._lock:
            self._flush_buffer()

    def load_agent_snapshot(
        self,
        agent_id: str,
        tool_counts: dict[str, int],
        *,
        session_id: str = DEFAULT_SESSION_ID,
    ) -> None:
        """Load precomputed tool counts for an agent (replay support)."""
        safe_agent_id = agent_id if isinstance(agent_id, str) and agent_id else GLOBAL_AGENT_ID
        safe_session_id = (
            session_id if isinstance(session_id, str) and session_id else DEFAULT_SESSION_ID
        )
        now = datetime.now(timezone.utc)
        with self._lock:
            for tool_name, count in tool_counts.items():
                if not isinstance(tool_name, str) or not isinstance(count, int):
                    continue
                if count <= 0:
                    continue
                key = (safe_agent_id, safe_session_id, tool_name)
                self._events[key] = [now for _ in range(count)]
