# Copyright (C) 2026 GoodData Corporation

"""Server-side session memory for get_rules delta mode."""

from __future__ import annotations

import hashlib
import os
import time
from collections import OrderedDict
from dataclasses import dataclass, field
from threading import RLock

from aida_mcp.tools.rules.formatting import RuleDigest


def _int_env(name: str, default: int) -> int:
    raw = os.getenv(name)
    if raw is None:
        return default
    try:
        value = int(raw)
    except ValueError:
        return default
    return value if value > 0 else default


MAX_SESSIONS = _int_env("AIDA_RULE_SESSIONS_MAX", 256)
MAX_DIGESTS_PER_SESSION = _int_env("AIDA_RULE_SESSION_MAX_DIGESTS", 2048)
SESSION_TTL_SECONDS = _int_env("AIDA_RULE_SESSION_TTL_SECONDS", 60 * 60 * 24)


@dataclass
class _SessionState:
    updated_at: float
    revision: int = 0
    digests_by_rule: OrderedDict[str, str] = field(default_factory=OrderedDict)


_LOCK = RLock()
_SESSIONS: OrderedDict[str, _SessionState] = OrderedDict()


def _prune_expired(now: float) -> None:
    expired_keys = [k for k, v in _SESSIONS.items() if now - v.updated_at > SESSION_TTL_SECONDS]
    for key in expired_keys:
        _SESSIONS.pop(key, None)


def _evict_if_needed() -> None:
    while len(_SESSIONS) > MAX_SESSIONS:
        _SESSIONS.popitem(last=False)


def _session_cursor(state: _SessionState) -> str:
    payload = "\n".join(f"{rule_id}:{sha256}" for rule_id, sha256 in state.digests_by_rule.items())
    fingerprint = hashlib.sha256(payload.encode("utf-8")).hexdigest()[:8]
    return f"{state.revision}.{fingerprint}"


def get_session_known_rule_digests(session_key: str) -> list[RuleDigest]:
    """Return known rule digests for a session key."""
    now = time.time()
    with _LOCK:
        _prune_expired(now)
        state = _SESSIONS.get(session_key)
        if state is None:
            return []
        state.updated_at = now
        _SESSIONS.move_to_end(session_key)
        return [{"rule_id": rid, "sha256": sha} for rid, sha in state.digests_by_rule.items()]


def get_session_cursor(session_key: str) -> str | None:
    """Return the current opaque cursor for a session key."""
    now = time.time()
    with _LOCK:
        _prune_expired(now)
        state = _SESSIONS.get(session_key)
        if state is None:
            return None
        state.updated_at = now
        _SESSIONS.move_to_end(session_key)
        return _session_cursor(state)


def merge_session_rule_digests(session_key: str, digests: list[RuleDigest]) -> tuple[int, str]:
    """Merge rule digests into a session and return resulting digest count and cursor."""
    now = time.time()
    with _LOCK:
        _prune_expired(now)
        state = _SESSIONS.get(session_key)
        if state is None:
            state = _SessionState(updated_at=now)
            _SESSIONS[session_key] = state
        state.updated_at = now
        _SESSIONS.move_to_end(session_key)

        changed = False
        for digest in digests:
            rule_id = digest.get("rule_id")
            sha256 = digest.get("sha256")
            if not rule_id or not sha256:
                continue
            existing = state.digests_by_rule.get(rule_id)
            if existing == sha256:
                state.digests_by_rule.move_to_end(rule_id)
                continue
            if rule_id in state.digests_by_rule:
                state.digests_by_rule.pop(rule_id)
            state.digests_by_rule[rule_id] = sha256
            changed = True

        while len(state.digests_by_rule) > MAX_DIGESTS_PER_SESSION:
            state.digests_by_rule.popitem(last=False)
            changed = True

        if changed:
            state.revision += 1

        _evict_if_needed()
        return len(state.digests_by_rule), _session_cursor(state)
