from __future__ import annotations

import json
from dataclasses import dataclass
from typing import Any, Awaitable, Callable, Dict, List, Mapping, Optional, Union

from ..hash import hash_input, hash_output
from ..receipt import create_receipt, encode_receipt_header
from ..sign import load_secret_key, public_key_from_secret, sign_receipt

try:
    from starlette.middleware.base import BaseHTTPMiddleware
    from starlette.requests import Request
    from starlette.responses import Response

    _STARLETTE_AVAILABLE = True
except Exception:  # pragma: no cover - optional dependency path
    BaseHTTPMiddleware = object  # type: ignore[assignment]
    Request = Any  # type: ignore[misc,assignment]
    Response = Any  # type: ignore[misc,assignment]
    _STARLETTE_AVAILABLE = False


PrincipalResolver = Callable[[Request], Mapping[str, Any]]
ScopeResolver = Callable[[Request], List[str]]
CostResolver = Callable[[Request, Response], Optional[Mapping[str, Any]]]
PersistCallback = Callable[[Mapping[str, Any]], Union[None, Awaitable[None]]]


@dataclass
class AARMiddlewareOptions:
    agent_id: str
    secret_key: Union[bytes, bytearray, str]
    agent_name: Optional[str] = None
    agent_version: Optional[str] = None
    header_name: Optional[str] = None
    principal_resolver: Optional[PrincipalResolver] = None
    scope_resolver: Optional[ScopeResolver] = None
    cost_resolver: Optional[CostResolver] = None
    persist: Optional[PersistCallback] = None


def _default_principal_resolver(request: Request) -> Dict[str, str]:
    wallet = request.headers.get("x-wallet")
    if wallet and wallet.strip():
        return {"id": wallet.strip(), "type": "user"}

    api_key = request.headers.get("x-api-key")
    if api_key and api_key.strip():
        return {"id": "apikey:%s..." % api_key.strip()[:8], "type": "service"}

    return {"id": "anonymous", "type": "other"}


def _default_scope_resolver(request: Request) -> List[str]:
    method = request.method.lower()
    path = request.url.path or "/"
    return ["%s:%s" % (method, path)]


def _parse_body(body: bytes, content_type: str) -> Any:
    if not body:
        return None

    if "application/json" in content_type:
        try:
            return json.loads(body.decode("utf-8"))
        except Exception:
            return body.decode("utf-8", errors="replace")

    try:
        return body.decode("utf-8")
    except Exception:
        return body.hex()


class AARMiddleware(BaseHTTPMiddleware):
    """FastAPI/Starlette middleware that emits signed AAR receipts on responses."""

    def __init__(
        self,
        app: Any,
        agent_id: str,
        secret_key: Union[bytes, bytearray, str],
        agent_name: Optional[str] = None,
        agent_version: Optional[str] = None,
        header_name: Optional[str] = None,
        principal_resolver: Optional[PrincipalResolver] = None,
        scope_resolver: Optional[ScopeResolver] = None,
        cost_resolver: Optional[CostResolver] = None,
        persist: Optional[PersistCallback] = None,
    ) -> None:
        if not _STARLETTE_AVAILABLE:  # pragma: no cover - optional dependency path
            raise ImportError("starlette/fastapi is required to use AARMiddleware")

        super().__init__(app)
        self._secret_key = load_secret_key(secret_key)
        self._public_key = public_key_from_secret(self._secret_key)
        self._agent_id = agent_id
        self._agent_name = agent_name
        self._agent_version = agent_version
        self._header_name = header_name or "X-AAR-Receipt"
        self._principal_resolver = principal_resolver or _default_principal_resolver
        self._scope_resolver = scope_resolver or _default_scope_resolver
        self._cost_resolver = cost_resolver or (lambda _req, _res: None)
        self._persist = persist

    async def dispatch(self, request: Request, call_next: Callable[[Request], Awaitable[Response]]) -> Response:
        body = await request.body()

        sent = False

        async def receive() -> Dict[str, Any]:
            nonlocal sent
            if sent:
                return {"type": "http.request", "body": b"", "more_body": False}
            sent = True
            return {"type": "http.request", "body": body, "more_body": False}

        req = Request(request.scope, receive)
        response = await call_next(req)

        output_chunks: List[bytes] = []
        async for chunk in response.body_iterator:
            output_chunks.append(chunk if isinstance(chunk, bytes) else bytes(chunk))
        output_body = b"".join(output_chunks)

        wrapped = Response(
            content=output_body,
            status_code=response.status_code,
            headers=dict(response.headers),
            media_type=response.media_type,
            background=response.background,
        )

        try:
            input_data = {
                "query": dict(req.query_params),
                "body": _parse_body(body, req.headers.get("content-type", "")),
                "method": req.method,
                "path": req.url.path,
            }

            cost = self._cost_resolver(req, wrapped) or {"amount": "0", "currency": "USD"}
            action_path = req.url.path or "/"

            unsigned = create_receipt(
                {
                    "agent": {
                        "id": self._agent_id,
                        "name": self._agent_name,
                        "version": self._agent_version,
                    },
                    "principal": dict(self._principal_resolver(req)),
                    "action": {
                        "type": "api.call",
                        "target": action_path,
                        "method": req.method.upper(),
                        "status": "success" if wrapped.status_code < 400 else "failure",
                    },
                    "scope": {"permissions": list(self._scope_resolver(req))},
                    "inputHash": hash_input(input_data),
                    "outputHash": hash_output(output_body),
                    "cost": dict(cost),
                }
            )

            receipt = sign_receipt(unsigned, self._secret_key)
            if self._header_name not in wrapped.headers:
                wrapped.headers[self._header_name] = encode_receipt_header(receipt)

            if self._persist is not None:
                maybe_awaitable = self._persist(receipt)
                if hasattr(maybe_awaitable, "__await__"):
                    await maybe_awaitable  # type: ignore[misc]
        except Exception:
            pass

        return wrapped


__all__ = ["AARMiddleware", "AARMiddlewareOptions"]
