from __future__ import annotations

import json
import os
import secrets
import time
from collections import OrderedDict
from typing import Any, Generic, TypeVar
from uuid import uuid4

import bcrypt
import httpx
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 AccountRow, 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())


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

_CACHE_MAX_SIZE = 1024
_CACHE_TTL_SECONDS = 300.0


class ApiKeyAuthMiddleware:
    def __init__(self, app: ASGIApp, db: RelayDB) -> None:
        self._app = app
        self._db = db
        self._key_cache: BoundedTTLCache[str, str] = 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", []))
        api_key = headers.get(b"x-api-key", b"").decode()

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

        email = self._key_cache.get(api_key)
        if email is None:
            account = await self._db.get_account_by_api_key(api_key)
            if account is None:
                response = JSONResponse({"error": "unauthorized"}, status_code=401)
                await response(scope, receive, send)
                return
            email = account.email
            self._key_cache.put(api_key, email)

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

        await self._app(scope, receive, send)


async def _authenticate_api_key(request: Request, db: RelayDB) -> AccountRow | None:
    api_key = request.headers.get("x-api-key", "")
    if not api_key:
        return None
    return await db.get_account_by_api_key(api_key)


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

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

        email = body.get("email", "")
        workos_user_id = body.get("workos_user_id", "")
        if not email or not workos_user_id:
            return JSONResponse({"error": "email and workos_user_id required"}, status_code=400)

        workos_api_key = os.environ.get("WORKOS_API_KEY", "")
        if not workos_api_key:
            return JSONResponse({"error": "server misconfigured"}, status_code=500)

        async with httpx.AsyncClient() as client:
            resp = await client.get(
                f"https://api.workos.com/user_management/users/{workos_user_id}",
                headers={"Authorization": f"Bearer {workos_api_key}"},
            )

        if resp.status_code != 200:
            return JSONResponse({"error": "invalid workos user"}, status_code=401)

        workos_data = resp.json()
        workos_email = workos_data.get("email", "")
        if workos_email != email:
            return JSONResponse({"error": "email mismatch"}, status_code=401)

        account = await db.get_account_by_email(email)
        if account is not None:
            return JSONResponse({"api_key": account.api_key, "email": email})

        api_key = secrets.token_urlsafe(32)
        await db.create_account(email, api_key)
        return JSONResponse({"api_key": api_key, "email": email})

    @mcp.custom_route("/auth/users", methods=["GET"])
    async def list_users(request: Request) -> Response:
        account = await _authenticate_api_key(request, db)
        if account is None:
            return JSONResponse({"error": "unauthorized"}, status_code=401)

        users = await db.list_users(account.email)
        return JSONResponse([
            {"user_id": u.user_id, "display_name": u.display_name}
            for u in users
        ])

    @mcp.custom_route("/auth/users", methods=["POST"])
    async def create_user(request: Request) -> Response:
        account = await _authenticate_api_key(request, db)
        if account is None:
            return JSONResponse({"error": "unauthorized"}, status_code=401)

        try:
            body: dict[str, Any] = json.loads(await request.body())
        except json.JSONDecodeError:
            return JSONResponse({"error": "invalid json"}, status_code=400)

        display_name = body.get("display_name", "")
        if not display_name:
            return JSONResponse({"error": "display_name is required"}, status_code=400)

        user_id = f"u-{uuid4().hex[:8]}"
        await db.create_user(user_id, account.email, display_name)
        return JSONResponse({"user_id": user_id, "display_name": display_name})

    @mcp.custom_route("/auth/users/{user_id}", methods=["DELETE"])
    async def delete_user(request: Request) -> Response:
        account = await _authenticate_api_key(request, db)
        if account is None:
            return JSONResponse({"error": "unauthorized"}, status_code=401)

        user_id = request.path_params["user_id"]
        user = await db.get_user(user_id)
        if user is None or user.account_email != account.email:
            return JSONResponse({"error": "not found"}, status_code=404)

        await db.delete_user(user_id)
        return JSONResponse({"ok": True})

    @mcp.custom_route("/auth/users/{user_id}/groups", methods=["GET"])
    async def get_user_groups(request: Request) -> Response:
        account = await _authenticate_api_key(request, db)
        if account is None:
            return JSONResponse({"error": "unauthorized"}, status_code=401)

        user_id = request.path_params["user_id"]
        user = await db.get_user(user_id)
        if user is None or user.account_email != account.email:
            return JSONResponse({"error": "not found"}, status_code=404)

        groups = await db.get_user_groups(user_id)
        return JSONResponse({"groups": groups})

    @mcp.custom_route("/auth/users/{user_id}/groups", methods=["POST"])
    async def add_user_to_group(request: Request) -> Response:
        account = await _authenticate_api_key(request, db)
        if account is None:
            return JSONResponse({"error": "unauthorized"}, status_code=401)

        user_id = request.path_params["user_id"]
        user = await db.get_user(user_id)
        if user is None or user.account_email != account.email:
            return JSONResponse({"error": "not found"}, status_code=404)

        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", "")
        password = body.get("password", "")

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

        if not bcrypt.checkpw(password.encode("utf-8"), stored_hash.encode("utf-8")):
            return JSONResponse({"error": "invalid password"}, status_code=401)

        await db.add_user_to_group(user_id, group_id)
        return JSONResponse({"ok": True})

    @mcp.custom_route("/auth/users/{user_id}/groups/{group_id}", methods=["DELETE"])
    async def remove_user_from_group(request: Request) -> Response:
        account = await _authenticate_api_key(request, db)
        if account is None:
            return JSONResponse({"error": "unauthorized"}, status_code=401)

        user_id = request.path_params["user_id"]
        user = await db.get_user(user_id)
        if user is None or user.account_email != account.email:
            return JSONResponse({"error": "not found"}, status_code=404)

        group_id = request.path_params["group_id"]
        await db.remove_user_from_group(user_id, group_id)
        return JSONResponse({"ok": True})
