"""Session management utilities for RapidKit projects."""

from __future__ import annotations

import base64
import json
from collections import defaultdict
import hashlib
import hmac
import os
import secrets
import time
from dataclasses import dataclass, field
from typing import Any, Dict, Mapping, MutableMapping, Optional


DEFAULTS: Dict[str, Any] = json.loads(
    """{{ module_defaults | default({}) | tojson(indent=2) }}"""
)

_FEATURE_FLAGS: tuple[str, ...] = (
    "signed_tokens",
    "pluggable_storage",
    "secure_cookies",
    "refresh_token_rotation",
)


def _env(name: str, fallback: str | None = None) -> str:
    if name:
        value = os.getenv(name)
        if value:
            return value
    return fallback or ""


@dataclass(slots=True)
class CookieSettings:
    """HTTP cookie configuration for session issuance."""

    name: str
    domain: Optional[str]
    secure: bool
    httponly: bool
    same_site: str


@dataclass(slots=True)
class SessionSettings:
    """Complete runtime settings for session management."""

    secret_key: bytes
    session_ttl_seconds: int
    refresh_ttl_seconds: int
    cookie: CookieSettings
    storage_backend: str


@dataclass(slots=True)
class SessionRecord:
    """Represents persisted session metadata."""

    session_id: str
    user_id: str
    issued_at: float
    expires_at: float
    payload: Dict[str, Any] = field(default_factory=dict)

    def is_expired(self, *, now: Optional[float] = None) -> bool:
        return (now or time.time()) >= self.expires_at


@dataclass(slots=True)
class RefreshRecord:
    """Stores refresh token state."""

    token: str
    session_id: str
    issued_at: float
    expires_at: float

    def is_expired(self, *, now: Optional[float] = None) -> bool:
        return (now or time.time()) >= self.expires_at


@dataclass(slots=True)
class SessionEnvelope:
    """Aggregated result from issuing or rotating a session."""

    session: SessionRecord
    token: str
    refresh_token: str
    cookie: Dict[str, Any]


class SessionSigner:
    """Handles signing and verifying session identifiers."""

    def __init__(self, secret: bytes) -> None:
        self._secret = secret

    def sign(self, session_id: str) -> str:
        mac = hmac.new(self._secret, session_id.encode("utf-8"), hashlib.sha256).digest()
        return base64.urlsafe_b64encode(mac).rstrip(b"=").decode("ascii")

    def verify(self, session_id: str, signature: str) -> bool:
        expected = self.sign(session_id)
        return hmac.compare_digest(expected, signature)


class InMemorySessionStore:
    """In-memory store suitable for development and unit tests."""

    def __init__(self) -> None:
        self._sessions: MutableMapping[str, SessionRecord] = {}

    def upsert(self, record: SessionRecord) -> None:
        self._sessions[record.session_id] = record

    def get(self, session_id: str) -> Optional[SessionRecord]:
        record = self._sessions.get(session_id)
        if record and record.is_expired():
            self.delete(session_id)
            return None
        return record

    def delete(self, session_id: str) -> None:
        self._sessions.pop(session_id, None)

    def prune(self) -> None:
        now = time.time()
        expired = [sid for sid, record in self._sessions.items() if record.is_expired(now=now)]
        for sid in expired:
            self.delete(sid)


class SessionRuntime:
    """High-level manager encapsulating session issuance and verification."""

    def __init__(self, settings: SessionSettings) -> None:
        self._settings = settings
        self._signer = SessionSigner(settings.secret_key)
        self._store = InMemorySessionStore()
        self._refresh_tokens: MutableMapping[str, RefreshRecord] = {}
        self._session_index: MutableMapping[str, set[str]] = defaultdict(set)

    @property
    def settings(self) -> SessionSettings:
        return self._settings

    def issue_session(
        self,
        user_id: str,
        *,
        payload: Optional[Mapping[str, Any]] = None,
        ttl_seconds: Optional[int] = None,
    ) -> SessionEnvelope:
        now = time.time()
        ttl = ttl_seconds or self._settings.session_ttl_seconds
        session_id = secrets.token_urlsafe(32)
        record = SessionRecord(
            session_id=session_id,
            user_id=user_id,
            issued_at=now,
            expires_at=now + ttl,
            payload=dict(payload or {}),
        )
        self._store.upsert(record)
        token = self._encode_token(session_id)
        refresh = self._issue_refresh_token(session_id, now)
        cookie = self._build_cookie(token, record.expires_at)
        return SessionEnvelope(session=record, token=token, refresh_token=refresh.token, cookie=cookie)

    def verify_session_token(self, token: str) -> SessionRecord:
        session_id = self._decode_token(token)
        record = self._store.get(session_id)
        if record is None:
            raise ValueError("Session token is invalid or has expired")
        return record

    def rotate_session(self, refresh_token: str) -> SessionEnvelope:
        refresh = self._refresh_tokens.get(refresh_token)
        if refresh is None or refresh.is_expired():
            raise ValueError("Refresh token is invalid or has expired")
        session = self._store.get(refresh.session_id)
        if session is None:
            raise ValueError("Associated session no longer exists")
        self.revoke_refresh_token(refresh_token)
        return self.issue_session(session.user_id, payload=session.payload)

    def revoke_session(self, session_id: str) -> None:
        self._store.delete(session_id)
        tokens = self._session_index.pop(session_id, set())
        for token in tokens:
            self._refresh_tokens.pop(token, None)

    def revoke_refresh_token(self, refresh_token: str) -> None:
        record = self._refresh_tokens.pop(refresh_token, None)
        if record:
            bucket = self._session_index.get(record.session_id)
            if bucket and refresh_token in bucket:
                bucket.remove(refresh_token)

    def _issue_refresh_token(self, session_id: str, issued_at: float) -> RefreshRecord:
        token = secrets.token_urlsafe(48)
        record = RefreshRecord(
            token=token,
            session_id=session_id,
            issued_at=issued_at,
            expires_at=issued_at + self._settings.refresh_ttl_seconds,
        )
        self._refresh_tokens[token] = record
        self._session_index[session_id].add(token)
        return record

    def _encode_token(self, session_id: str) -> str:
        signature = self._signer.sign(session_id)
        return f"{session_id}.{signature}"

    def _decode_token(self, token: str) -> str:
        try:
            session_id, signature = token.split(".", 1)
        except ValueError as exc:  # pragma: no cover - defensive
            raise ValueError("Invalid session token format") from exc
        if not self._signer.verify(session_id, signature):
            raise ValueError("Invalid session token signature")
        return session_id

    def _build_cookie(self, value: str, expires_at: float) -> Dict[str, Any]:
        return {
            "name": self._settings.cookie.name,
            "value": value,
            "domain": self._settings.cookie.domain,
            "secure": self._settings.cookie.secure,
            "httponly": self._settings.cookie.httponly,
            "same_site": self._settings.cookie.same_site,
            "expires": int(expires_at),
        }

    def metadata(self) -> Dict[str, Any]:
        return describe_session(self._settings)


def load_session_settings(overrides: Optional[Mapping[str, Any]] = None) -> SessionSettings:
    """Load session settings, applying optional overrides."""

    config: Dict[str, Any] = dict(DEFAULTS)
    if overrides:
        config.update(overrides)

    secret_env = str(config.get("secret_key_env", ""))
    secret_value = _env(secret_env, str(config.get("secret_key", "")))
    if not secret_value:
        raise RuntimeError(
            "Session secret missing. Set the environment variable defined by 'secret_key_env'."
        )

    cookie = CookieSettings(
        name=str(config.get("cookie_name", "rapidkit_session")),
        domain=config.get("cookie_domain"),
        secure=bool(config.get("cookie_secure", True)),
        httponly=bool(config.get("cookie_httponly", True)),
        same_site=str(config.get("cookie_same_site", "lax")).lower(),
    )

    return SessionSettings(
        secret_key=secret_value.encode("utf-8"),
        session_ttl_seconds=int(config.get("session_ttl_seconds", 30 * 24 * 60 * 60)),
        refresh_ttl_seconds=int(config.get("refresh_ttl_seconds", 180 * 24 * 60 * 60)),
        cookie=cookie,
        storage_backend=str(config.get("storage_backend", "memory")),
    )


def describe_session(settings: Optional[SessionSettings] = None) -> Dict[str, Any]:
    """Return a metadata payload describing the session runtime."""

    config = settings or load_session_settings()
    return {
        "module": "{{ module_name | default('session') }}",
        "session_ttl_seconds": config.session_ttl_seconds,
        "refresh_ttl_seconds": config.refresh_ttl_seconds,
        "storage_backend": config.storage_backend,
        "cookie": {
            "name": config.cookie.name,
            "domain": config.cookie.domain,
            "secure": config.cookie.secure,
            "httponly": config.cookie.httponly,
            "same_site": config.cookie.same_site,
        },
        "features": list_session_features(),
        "supports_refresh_tokens": True,
    }


def list_session_features() -> list[str]:
    """Enumerate capabilities surfaced by the session module."""

    return list(_FEATURE_FLAGS)


def get_session_metadata(settings: Optional[SessionSettings] = None) -> Dict[str, Any]:
    """Convenience wrapper mirroring the runtime metadata signature."""

    return describe_session(settings)


__all__ = [
    "CookieSettings",
    "SessionSettings",
    "SessionRecord",
    "RefreshRecord",
    "SessionEnvelope",
    "SessionRuntime",
    "load_session_settings",
    "describe_session",
    "list_session_features",
    "get_session_metadata",
]
