from __future__ import annotations

import asyncio
import json
from typing import Any, 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

PrincipalResolver = Callable[[Any], Mapping[str, Any]]
ScopeResolver = Callable[[Any], List[str]]
CostResolver = Callable[[Any, Any], Optional[Mapping[str, Any]]]
PersistCallback = Callable[[Mapping[str, Any]], Any]


def _default_principal_resolver(request: Any) -> 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: Any) -> List[str]:
    return ["%s:%s" % (request.method.lower(), request.path or "/")]


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


def install_aar_middleware(
    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:
    try:
        from flask import g, request
    except Exception as exc:  # pragma: no cover - optional dependency path
        raise ImportError("flask is required to use install_aar_middleware") from exc

    sk = load_secret_key(secret_key)
    _pk = public_key_from_secret(sk)
    _ = _pk

    resolved_header = header_name or "X-AAR-Receipt"
    resolve_principal = principal_resolver or _default_principal_resolver
    resolve_scope = scope_resolver or _default_scope_resolver
    resolve_cost = cost_resolver or (lambda _req, _res: None)

    @app.before_request
    def _capture_body() -> None:
        g._aar_request_body = request.get_data(cache=True)

    @app.after_request
    def _attach_receipt(response: Any) -> Any:
        try:
            request_body = getattr(g, "_aar_request_body", b"")
            output_body = response.get_data() or b""

            input_data = {
                "query": request.args.to_dict(flat=True),
                "body": _parse_body(request_body, request.content_type or ""),
                "method": request.method,
                "path": request.path,
            }

            cost = resolve_cost(request, response) or {"amount": "0", "currency": "USD"}

            unsigned = create_receipt(
                {
                    "agent": {
                        "id": agent_id,
                        "name": agent_name,
                        "version": agent_version,
                    },
                    "principal": dict(resolve_principal(request)),
                    "action": {
                        "type": "api.call",
                        "target": request.path or "/",
                        "method": request.method.upper(),
                        "status": "success" if response.status_code < 400 else "failure",
                    },
                    "scope": {"permissions": list(resolve_scope(request))},
                    "inputHash": hash_input(input_data),
                    "outputHash": hash_output(output_body),
                    "cost": dict(cost),
                }
            )

            receipt = sign_receipt(unsigned, sk)
            if resolved_header not in response.headers:
                response.headers[resolved_header] = encode_receipt_header(receipt)

            if persist is not None:
                result = persist(receipt)
                if hasattr(result, "__await__"):
                    try:
                        loop = asyncio.get_running_loop()
                        loop.create_task(result)
                    except RuntimeError:
                        pass
        except Exception:
            pass

        return response


aar_middleware = install_aar_middleware

__all__ = ["install_aar_middleware", "aar_middleware"]
