from __future__ import annotations

import json
import secrets
import time
from collections import OrderedDict
from dataclasses import dataclass
from typing import Any, Generic, TypeVar

import bcrypt
from fastmcp import FastMCP
from starlette.requests import Request
from starlette.responses import JSONResponse, Response
from starlette.types import ASGIApp, Receive, Scope, Send

from relay.server.db import RelayDB

K = TypeVar("K")
V = TypeVar("V")


class BoundedTTLCache(Generic[K, V]):
    def __init__(self, maxsize: int, ttl: float) -> None:
        self._maxsize = maxsize
        self._ttl = ttl
        self._data: OrderedDict[K, tuple[V, float]] = OrderedDict()

    def get(self, key: K) -> V | None:
        entry = self._data.get(key)
        if entry is None:
            return None
        value, ts = entry
        if time.monotonic() - ts > self._ttl:
            del self._data[key]
            return None
        self._data.move_to_end(key)
        return value

    def put(self, key: K, value: V) -> None:
        if key in self._data:
            self._data.move_to_end(key)
            self._data[key] = (value, time.monotonic())
            return
        if len(self._data) >= self._maxsize:
            self._data.popitem(last=False)
        self._data[key] = (value, time.monotonic())


@dataclass(frozen=True, slots=True)
class AuthInfo:
    peer_id: str
    authorized_groups: list[str]
    hostname: str
    working_directory: str


class AuthStore:
    def __init__(self) -> None:
        self._sessions: dict[str, AuthInfo] = {}

    def store(self, session_id: str, info: AuthInfo) -> None:
        self._sessions[session_id] = info

    def get(self, session_id: str) -> AuthInfo:
        info = self._sessions.get(session_id)
        if info is None:
            raise ValueError(f"No auth context for session: {session_id}")
        return info

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


PROTECTED_PREFIXES = ("/mcp", "/notifications")
PUBLIC_PREFIXES = ("/auth",)


_CACHE_MAX_SIZE = 1024
_CACHE_TTL_SECONDS = 300.0


class GroupAuthMiddleware:
    def __init__(self, app: ASGIApp, db: RelayDB, auth_store: AuthStore) -> None:
        self._app = app
        self._db = db
        self._auth_store = auth_store
        self._verified_cache: BoundedTTLCache[tuple[str, str], bool] = BoundedTTLCache(
            maxsize=_CACHE_MAX_SIZE, ttl=_CACHE_TTL_SECONDS
        )

    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
        if scope["type"] != "http":
            await self._app(scope, receive, send)
            return

        path: str = scope.get("path", "")

        if any(path.startswith(p) for p in PUBLIC_PREFIXES):
            await self._app(scope, receive, send)
            return

        if not any(path.startswith(p) for p in PROTECTED_PREFIXES):
            await self._app(scope, receive, send)
            return

        headers = dict(scope.get("headers", []))
        peer_id = headers.get(b"x-peer-id", b"").decode()
        groups_raw = headers.get(b"x-groups", b"").decode()
        hostname = headers.get(b"x-hostname", b"unknown").decode()
        working_directory = headers.get(b"x-working-directory", b"unknown").decode()

        authorized_groups: list[str] = []
        seen_groups: set[str] = set()
        if groups_raw:
            for pair in groups_raw.split(","):
                parts = pair.strip().split(":", 1)
                if len(parts) != 2:
                    continue
                group_id, password = parts
                if group_id in seen_groups:
                    continue
                seen_groups.add(group_id)
                authorized = await self._check_group(group_id, password)
                if authorized:
                    authorized_groups.append(group_id)

        if not authorized_groups:
            response = JSONResponse(
                {"error": "unauthorized"}, status_code=401
            )
            await response(scope, receive, send)
            return

        auth_info = AuthInfo(
            peer_id=peer_id,
            authorized_groups=authorized_groups,
            hostname=hostname,
            working_directory=working_directory,
        )

        mcp_session_id = headers.get(b"mcp-session-id", b"").decode()
        if mcp_session_id:
            self._auth_store.store(mcp_session_id, auth_info)

        scope["state"] = scope.get("state", {})
        scope["state"]["auth_info"] = auth_info

        await self._app(scope, receive, send)

    async def _check_group(self, group_id: str, password: str) -> bool:
        cache_key = (group_id, password)
        cached = self._verified_cache.get(cache_key)
        if cached is not None:
            return cached

        stored_hash = await self._db.get_group_password_hash(group_id)
        if stored_hash is None:
            self._verified_cache.put(cache_key, False)
            return False

        valid = bcrypt.checkpw(
            password.encode("utf-8"), stored_hash.encode("utf-8")
        )
        self._verified_cache.put(cache_key, valid)
        return valid


def register_auth_routes(mcp: FastMCP, db: RelayDB) -> None:
    @mcp.custom_route("/auth/groups", methods=["GET"])
    async def list_groups(request: Request) -> Response:
        groups = await db.list_groups()
        return JSONResponse(groups)

    @mcp.custom_route("/auth/groups", methods=["POST"])
    async def create_group(request: Request) -> Response:
        try:
            body: dict[str, Any] = json.loads(await request.body())
        except json.JSONDecodeError:
            return JSONResponse({"error": "invalid json"}, status_code=400)
        group_id = body.get("group_id", "")
        if not group_id:
            return JSONResponse(
                {"error": "group_id is required"}, status_code=400
            )

        if await db.group_exists(group_id):
            return JSONResponse(
                {"error": "group already exists"}, status_code=409
            )

        password = secrets.token_urlsafe(6)
        password_hash = bcrypt.hashpw(
            password.encode("utf-8"), bcrypt.gensalt()
        ).decode("utf-8")

        await db.create_group(group_id, password_hash)
        return JSONResponse({"group_id": group_id, "password": password})

    @mcp.custom_route("/auth/groups/{group_id}/verify", methods=["POST"])
    async def verify_group(request: Request) -> Response:
        group_id = request.path_params["group_id"]
        try:
            body: dict[str, Any] = json.loads(await request.body())
        except json.JSONDecodeError:
            return JSONResponse({"error": "invalid json"}, status_code=400)
        password = body.get("password", "")

        stored_hash = await db.get_group_password_hash(group_id)
        if stored_hash is None:
            return JSONResponse(
                {"ok": False, "error": "group not found"}, status_code=404
            )

        ok = bcrypt.checkpw(
            password.encode("utf-8"), stored_hash.encode("utf-8")
        )
        return JSONResponse({"ok": ok})
