from __future__ import annotations

import json
import secrets
from contextvars import ContextVar
from typing import Any

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

_peer_id_var: ContextVar[str] = ContextVar("auth_peer_id")
_authorized_groups_var: ContextVar[list[str]] = ContextVar("auth_authorized_groups")


def set_auth_context(peer_id: str, authorized_groups: list[str]) -> None:
    _peer_id_var.set(peer_id)
    _authorized_groups_var.set(authorized_groups)


def get_auth_context() -> tuple[str, list[str]]:
    try:
        return _peer_id_var.get(), _authorized_groups_var.get()
    except LookupError:
        raise ValueError("Auth context not available")


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


class GroupAuthMiddleware:
    def __init__(self, app: ASGIApp, db: RelayDB) -> None:
        self._app = app
        self._db = db
        self._verified_cache: dict[tuple[str, str], bool] = {}

    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()

        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

        scope["state"] = scope.get("state", {})
        scope["state"]["peer_id"] = peer_id
        scope["state"]["authorized_groups"] = authorized_groups
        set_auth_context(peer_id, authorized_groups)

        await self._app(scope, receive, send)

    async def _check_group(self, group_id: str, password: str) -> bool:
        cache_key = (group_id, password)
        if cache_key in self._verified_cache:
            return self._verified_cache[cache_key]

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

        valid = bcrypt.checkpw(
            password.encode("utf-8"), stored_hash.encode("utf-8")
        )
        self._verified_cache[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})
