"""Opaque server-side OAuth state token helpers."""

from __future__ import annotations

import json
import os
import secrets
import threading
import time
from contextlib import contextmanager
from datetime import UTC, datetime
from typing import TYPE_CHECKING, Any, cast
from uuid import uuid4

from mindroom.file_locks import advisory_file_lock
from mindroom.logging_config import get_logger
from mindroom.oauth.providers import OAuthProviderError

if TYPE_CHECKING:
    from collections.abc import Iterator
    from pathlib import Path

    from mindroom.constants import RuntimePaths

_oauth_state_lock = threading.Lock()
_OAUTH_STATE_DIR_NAME = "oauth_state"
_OAUTH_STATE_FILE_NAME = "oauth_state.json"
logger = get_logger(__name__)


def _state_file(runtime_paths: RuntimePaths) -> Path:
    return runtime_paths.storage_root / _OAUTH_STATE_DIR_NAME / _OAUTH_STATE_FILE_NAME


def _state_lock_file(runtime_paths: RuntimePaths) -> Path:
    return runtime_paths.storage_root / _OAUTH_STATE_DIR_NAME / f"{_OAUTH_STATE_FILE_NAME}.lock"


@contextmanager
def _locked_state_store(
    runtime_paths: RuntimePaths,
    *,
    now: float,
    save_on_exit: bool = True,
) -> Iterator[dict[str, dict[str, Any]]]:
    with _oauth_state_lock, advisory_file_lock(_state_lock_file(runtime_paths)):
        states, load_failed = _load_state_store(runtime_paths, now=now)
        initial_states = dict(states)
        yield states
        if save_on_exit and (not load_failed or states != initial_states):
            _save_state_store(runtime_paths, states)


def _corrupt_state_file(path: Path) -> Path:
    timestamp = datetime.now(UTC).isoformat(timespec="seconds")
    corrupt_path = path.with_name(f"{path.name}.corrupt-{timestamp}")
    if corrupt_path.exists():
        corrupt_path = path.with_name(f"{path.name}.corrupt-{timestamp}-{uuid4().hex}")
    path.replace(corrupt_path)
    return corrupt_path


def _load_state_store(runtime_paths: RuntimePaths, *, now: float) -> tuple[dict[str, dict[str, Any]], bool]:
    path = _state_file(runtime_paths)
    if not path.exists():
        return {}, False
    try:
        raw = json.loads(path.read_text(encoding="utf-8"))
    except json.JSONDecodeError as exc:
        corrupt_path = _corrupt_state_file(path)
        logger.warning(
            "oauth_state_store_corrupt",
            path=str(path),
            corrupt_path=str(corrupt_path),
            error=str(exc),
        )
        return {}, True
    except OSError:
        return {}, True
    if not isinstance(raw, dict):
        return {}, False
    states = raw.get("states")
    if not isinstance(states, dict):
        return {}, False
    pruned: dict[str, dict[str, Any]] = {}
    for token, record in states.items():
        if not isinstance(token, str) or not isinstance(record, dict):
            continue
        expires_at = record.get("exp")
        if isinstance(expires_at, int | float) and expires_at > now:
            pruned[token] = record
    return pruned, False


def _save_state_store(runtime_paths: RuntimePaths, states: dict[str, dict[str, Any]]) -> None:
    path = _state_file(runtime_paths)
    path.parent.mkdir(parents=True, exist_ok=True)
    tmp_path = path.with_name(f"{path.name}.tmp-{os.getpid()}-{uuid4().hex}")
    tmp_path.write_text(json.dumps({"states": states}, sort_keys=True, separators=(",", ":")), encoding="utf-8")
    tmp_path.replace(path)


def issue_opaque_oauth_state(
    runtime_paths: RuntimePaths,
    *,
    kind: str,
    ttl_seconds: int,
    data: dict[str, Any],
) -> str:
    """Return one opaque, time-limited OAuth state token."""
    now = time.time()
    token = secrets.token_urlsafe(32)
    record = {
        "kind": kind,
        "iat": now,
        "exp": now + ttl_seconds,
        "data": dict(data),
    }
    with _locked_state_store(runtime_paths, now=now) as states:
        states[token] = record
    return token


def _validated_oauth_state_data(record: object, *, kind: str, now: float) -> dict[str, Any]:
    if not isinstance(record, dict):
        msg = "OAuth state is invalid or expired"
        raise OAuthProviderError(msg)
    state_record = cast("dict[str, Any]", record)
    if state_record.get("kind") != kind:
        msg = "OAuth state does not match this integration"
        raise OAuthProviderError(msg)
    expires_at = state_record.get("exp")
    if not isinstance(expires_at, int | float) or expires_at <= now:
        msg = "OAuth state is invalid or expired"
        raise OAuthProviderError(msg)
    data = state_record.get("data")
    if not isinstance(data, dict):
        msg = "OAuth state is invalid or expired"
        raise OAuthProviderError(msg)
    return data


def read_opaque_oauth_state(
    runtime_paths: RuntimePaths,
    *,
    kind: str,
    token: str,
) -> dict[str, Any]:
    """Return one server-side OAuth state payload without consuming it."""
    now = time.time()
    with _locked_state_store(runtime_paths, now=now, save_on_exit=False) as states:
        record = states.get(token)

    return _validated_oauth_state_data(record, kind=kind, now=now)


def consume_opaque_oauth_state(
    runtime_paths: RuntimePaths,
    *,
    kind: str,
    token: str,
) -> dict[str, Any]:
    """Return and remove one server-side OAuth state payload."""
    now = time.time()
    with _locked_state_store(runtime_paths, now=now) as states:
        record = states.pop(token, None)
    return _validated_oauth_state_data(record, kind=kind, now=now)
