"""
Animica RPC — JSON-RPC 2.0 Dispatcher
=====================================

Features
--------
• Full JSON-RPC 2.0 compliance: single & batch, named & positional params, notifications.
• Structured error mapping (standard codes + app-defined via rpc.errors).
• Async-aware execution; works with FastAPI (router provided).
• Safe arg binding with optional context injection ("ctx"/"context"/"request").
• Deterministic responses: {"jsonrpc":"2.0", "id":..., "result":...} or {"error":...}.

This module is framework-light and can be reused outside FastAPI. For HTTP usage,
import `router` into rpc/server.py and mount at `/rpc`.
"""

from __future__ import annotations

import asyncio
import base64
import binascii
import difflib
import inspect
import json
import logging
import time
import uuid
from dataclasses import dataclass
from pathlib import Path
from typing import (Any, Awaitable, Callable, Dict, Iterable, List, Optional,
                    Tuple, Union)

from fastapi import APIRouter, Request, Response

try:
    # Prefer our shared error types if available
    from .errors import (InternalError, InvalidParams, InvalidRequest,
                         MethodNotFound, RpcError, to_error)
except Exception:  # pragma: no cover - fallback if errors module not ready

    class JsonRpcError(Exception):  # type: ignore
        code: int = -32000
        message: str = "Server error"
        data: Any = None

        def __init__(
            self,
            message: Optional[str] = None,
            *,
            code: Optional[int] = None,
            data: Any = None,
        ):
            if message is not None:
                self.message = message
            if code is not None:
                self.code = code
            self.data = data
            super().__init__(self.message)

    class InvalidRequest(JsonRpcError):  # type: ignore
        code = -32600
        message = "Invalid Request"

    class MethodNotFound(JsonRpcError):  # type: ignore
        code = -32601
        message = "Method not found"

    class InvalidParams(JsonRpcError):  # type: ignore
        code = -32602
        message = "Invalid params"

    class InternalError(JsonRpcError):  # type: ignore
        code = -32603
        message = "Internal error"

    class RpcError(JsonRpcError):  # type: ignore
        def __init__(self, message: str, code: int, data: Any | None = None):
            self.data = data
            self.code = code
            self.message = message
            super().__init__(message)

        def to_dict(self) -> dict[str, Any]:
            base = {"code": int(self.code), "message": self.message}
            if self.data is not None:
                base["data"] = self.data
            return base

    def to_error(exc: Exception) -> RpcError:
        if isinstance(exc, RpcError):
            return exc
        if isinstance(exc, JsonRpcError):
            return RpcError(
                getattr(exc, "message", str(exc)),
                getattr(exc, "code", -32000),
                getattr(exc, "data", None),
            )
        if isinstance(exc, (ValueError, TypeError)):
            return InvalidParams(str(exc))
        return InternalError(str(exc))


from rpc import version as rpc_version
from rpc.methods import ensure_loaded
from rpc.access_policy import get_active_policy

log = logging.getLogger(__name__)


def _extract_peer_client(request: Request) -> tuple[str, int] | None:
    """Best-effort transport peer extraction without trusting proxy headers."""
    candidates: list[tuple[Any, Any]] = []
    try:
        scope_client = request.scope.get("client") if isinstance(request.scope, dict) else None
        if isinstance(scope_client, tuple) and len(scope_client) >= 2:
            candidates.append((scope_client[0], scope_client[1]))
    except Exception:
        pass
    try:
        rc = request.client
        if rc is not None:
            candidates.append((getattr(rc, "host", None), getattr(rc, "port", None)))
    except Exception:
        pass

    for host_raw, port_raw in candidates:
        host = str(host_raw or "").strip()
        if not host:
            continue
        try:
            port = int(port_raw)
        except Exception:
            port = 0
        return (host, port)
    return None

Json = Dict[str, Any]
Params = Union[List[Any], Dict[str, Any]]
CallableLike = Union[Callable[..., Any], Callable[..., Awaitable[Any]]]


# --------------------------------------------------------------------------------------
# Context
# --------------------------------------------------------------------------------------


@dataclass
class Context:
    """
    Minimal per-request context passed to methods when they accept an arg named
    'ctx' or 'context' (or 'request' for FastAPI Request).

    Extend this as needed (e.g., auth/session, rate limits).
    """

    request: Optional[Request]
    received_at_ms: int
    client: Optional[Tuple[str, int]]
    headers: Dict[str, str]


def _now_ms() -> int:
    return int(time.time() * 1000)


# --------------------------------------------------------------------------------------
# Method registry
# --------------------------------------------------------------------------------------


class MethodRegistry:
    """
    Name → callable registry with decorator sugar.

    Methods can be sync or async. They may declare an argument named 'ctx',
    'context', or 'request' to receive the Context/Request object.
    """

    def __init__(self) -> None:
        self._methods: Dict[str, CallableLike] = {}

    def method(self, name: str) -> Callable[[CallableLike], CallableLike]:
        def deco(fn: CallableLike) -> CallableLike:
            if not isinstance(name, str) or not name:
                raise ValueError("Method name must be non-empty string")
            if name in self._methods:
                raise ValueError(f"Method already registered: {name}")
            self._methods[name] = fn
            log.debug("JSON-RPC register %s → %s.%s", name, fn.__module__, fn.__name__)
            return fn

        return deco

    def register(self, name: str, fn: CallableLike) -> None:
        self.method(name)(fn)

    def get(self, name: str) -> CallableLike:
        fn = self._methods.get(name)
        if fn is None:
            raise MethodNotFound(name)
        return fn

    @property
    def names(self) -> List[str]:
        return sorted(self._methods.keys())


registry = MethodRegistry()


def _sync_with_methods_registry() -> None:
    """Register handlers from rpc.methods into the local registry.

    The rpc.methods registry is the canonical source for namespaced JSON-RPC
    handlers (chain.*, tx.*, state.*, …). The lightweight dispatcher in this
    module keeps its own registry, so we mirror the entries at import time to
    avoid "Method not found" errors and to keep a single source of truth for
    method bindings.
    """

    try:
        from rpc import methods as method_registry

        method_registry.ensure_loaded()
        for name, spec in method_registry.get_registry().items():
            try:
                registry.register(name, spec.func)
                for alias in getattr(spec, "aliases", ()):
                    registry.register(alias, spec.func)
            except Exception:
                # Avoid hard-failing import if a method is already present or
                # missing; the dispatcher will surface a MethodNotFound at call
                # time if nothing is registered.
                continue
    except Exception:
        log.exception("Failed to sync rpc.methods registry into jsonrpc dispatcher")


_sync_with_methods_registry()

# Cached OpenRPC doc (if available)
_OPENRPC_CACHE_PATH: Path | None = None
_OPENRPC_CACHE_DOC: Dict[str, Any] | None = None
_OPENRPC_CACHE_MTIME: float | None = None


# --------------------------------------------------------------------------------------
# OpenRPC discovery helper
# --------------------------------------------------------------------------------------


def _load_openrpc_document() -> Dict[str, Any]:
    """
    Return the OpenRPC document for this server.

    Prefer loading from the same path served by /openrpc.json so rpc.discover
    mirrors the HTTP endpoint. Falls back to a synthesized document derived from
    the registered method names if the file is missing or invalid.
    """

    global _OPENRPC_CACHE_DOC, _OPENRPC_CACHE_PATH, _OPENRPC_CACHE_MTIME

    path: Path | None = None
    try:
        # Import lazily to avoid circular imports at module load time.
        from rpc.openrpc_mount import _resolve_openrpc_path  # type: ignore

        path = Path(_resolve_openrpc_path()).expanduser()
    except Exception:
        path = None

    if path and path.exists():
        try:
            mtime = path.stat().st_mtime
            if (
                _OPENRPC_CACHE_DOC is not None
                and _OPENRPC_CACHE_PATH == path
                and _OPENRPC_CACHE_MTIME == mtime
            ):
                return _OPENRPC_CACHE_DOC

            with path.open("r", encoding="utf-8") as f:
                doc = json.load(f)

            _OPENRPC_CACHE_PATH = path
            _OPENRPC_CACHE_MTIME = mtime
            _OPENRPC_CACHE_DOC = doc
        except Exception as e:  # pragma: no cover - best-effort cache
            log.warning("Failed to load OpenRPC document from %s: %s", path, e)

    if _OPENRPC_CACHE_DOC is None:
        doc: Dict[str, Any] = {
            "openrpc": "1.2.6",
            "info": {
                "title": "Animica RPC",
                "version": getattr(rpc_version, "__version__", "dev"),
            },
            "methods": [{"name": name} for name in registry.names],
        }
    else:
        doc = dict(_OPENRPC_CACHE_DOC)

    methods_list = list(doc.get("methods", []))
    known_names = set()
    for m in methods_list:
        if isinstance(m, str):
            known_names.add(m)
        elif isinstance(m, dict) and "name" in m:
            known_names.add(str(m["name"]))

    for name in registry.names:
        if name not in known_names:
            methods_list.append({"name": name})

    doc["methods"] = methods_list
    _OPENRPC_CACHE_DOC = doc
    return doc


# --------------------------------------------------------------------------------------
# Error shaping
# --------------------------------------------------------------------------------------


def _error_obj(exc: Exception) -> Json:
    """
    Convert any exception into a JSON-RPC error object.

    If `exc` is a JsonRpcError (or subclass), we preserve `.code`, `.message`,
    and include `.data` if present. Otherwise:
      • ValueError/TypeError → InvalidParams
      • Otherwise → Server error (-32000) with message
    """
    rpc_err = to_error(exc)
    code = getattr(rpc_err, "code", -32603)
    message = getattr(rpc_err, "message", "Internal error")
    data = getattr(rpc_err, "data", None)
    original_message = message

    # Normalize canonical JSON-RPC messages
    if code == -32602:
        if message == "Invalid params":
            message = "Invalid params"
    elif code == -32601:
        message = "Method not found"
    elif code == -32603:
        # Preserve custom details for internal errors when available
        if original_message != "Internal error":
            message = original_message
            data = data or {"detail": original_message}
        else:
            message = "Internal error"

    if hasattr(rpc_err, "to_dict"):
        base = rpc_err.to_dict()  # type: ignore[assignment]
        base["code"] = code
        base["message"] = message
        if data is not None:
            base["data"] = data
        return base

    # Fallback for unexpected shapes
    return {
        "code": code,
        "message": message,
        "data": data,
    }


_REDACT_KEYS = {
    "password",
    "secret",
    "seed",
    "mnemonic",
    "privkey",
    "private_key",
    "privateKey",
    "token",
    "auth",
    "signature",
    "sig",
    "rawTx",
}


def _redact_params(params: Any) -> Any:
    if isinstance(params, dict):
        redacted: dict[str, Any] = {}
        for key, value in params.items():
            if key in _REDACT_KEYS:
                redacted[key] = "<redacted>"
            else:
                redacted[key] = _redact_params(value)
        return redacted
    if isinstance(params, list):
        return [_redact_params(item) for item in params]
    return params


# --------------------------------------------------------------------------------------
# Arg binding & execution
# --------------------------------------------------------------------------------------


def _bind_call_args(
    fn: CallableLike, params: Optional[Params], ctx: Context
) -> Tuple[List[Any], Dict[str, Any]]:
    """
    Bind positional/named params to `fn` using its signature. Injects context into
    parameters named 'ctx'/'context'/'request' if not provided by caller.
    """
    sig = inspect.signature(fn)

    # Normalize params
    if params is None:
        args_obj: Params = []
    else:
        args_obj = params
        if not isinstance(args_obj, (list, dict)):
            args_obj = [args_obj]

    try:
        if isinstance(args_obj, list):
            bound = sig.bind_partial(*args_obj)  # allow extra defaults
        elif isinstance(args_obj, dict):
            if "address" in sig.parameters and "address" not in args_obj:
                if "addr" in args_obj:
                    addr_value = args_obj.get("addr")
                    args_obj = {k: v for k, v in args_obj.items() if k != "addr"}
                    args_obj["address"] = addr_value
            bound = sig.bind_partial(**args_obj)
        else:
            raise InvalidParams("params must be array or object")
    except TypeError as e:
        # Signature mismatch (wrong arity/unknown kw)
        raise InvalidParams(str(e))

    # Optional context injection
    for want in ("ctx", "context"):
        if want in sig.parameters and want not in bound.arguments:
            bound.arguments[want] = ctx
    if "request" in sig.parameters and "request" not in bound.arguments:
        bound.arguments["request"] = ctx.request

    missing: list[str] = []
    for name, param in sig.parameters.items():
        if name in bound.arguments:
            continue
        if param.kind in (inspect.Parameter.VAR_POSITIONAL, inspect.Parameter.VAR_KEYWORD):
            continue
        if name in {"ctx", "context", "request"}:
            continue
        if param.default is inspect.Parameter.empty:
            missing.append(name)
    if missing:
        raise InvalidParams(f"missing required params: {', '.join(missing)}")

    # Preserve positional arguments when present (e.g., varargs handlers) while
    # still honoring keyword-only bindings.
    args = list(bound.args)
    kwargs = dict(bound.kwargs)
    return args, kwargs


async def _maybe_await(x: Any) -> Any:
    if inspect.isawaitable(x):
        return await x  # type: ignore[no-any-return]
    return x


async def _invoke_handler(fn: Callable[..., Any], args: List[Any], kwargs: Dict[str, Any]) -> Any:
    # Async handlers run on the event loop. Sync handlers go through
    # asyncio.to_thread so blocking I/O (SQLite reads, template builds,
    # signature verification) doesn't stall every other in-flight RPC —
    # one slow handler used to freeze the entire event loop, which would
    # cascade into wedged-server territory under load.
    if inspect.iscoroutinefunction(fn):
        return await fn(*args, **kwargs)
    return await asyncio.to_thread(fn, *args, **kwargs)


# --------------------------------------------------------------------------------------
# Core dispatch
# --------------------------------------------------------------------------------------


def _validate_id(id_val: Any) -> Any:
    # Spec allows string, number, or null for id
    if id_val is None or isinstance(id_val, (str, int, float)):
        return id_val
    # If id is invalid type, the request itself is invalid
    raise InvalidRequest("id must be string, number, or null")


def _validate_request_obj(obj: Json) -> Tuple[str, Optional[Params], Any]:
    """
    Validate base request object; returns (method, params, id).
    Raises InvalidRequest on structural errors. Does NOT validate method existence.
    """
    if not isinstance(obj, dict):
        raise InvalidRequest("Request must be an object")

    if obj.get("jsonrpc") != "2.0":
        raise InvalidRequest("jsonrpc must be '2.0'")

    method = obj.get("method")
    if not isinstance(method, str) or not method:
        raise InvalidRequest("method must be a non-empty string")

    params: Optional[Params] = obj.get("params")
    if params is None:
        params = []
    elif isinstance(params, str):
        raw = params.strip()
        if raw:
            try:
                params = json.loads(raw)
            except Exception:
                try:
                    decoded = base64.b64decode(raw, validate=True)
                    params = json.loads(decoded.decode("utf-8"))
                except Exception:
                    try:
                        hex_raw = raw[2:] if raw.lower().startswith("0x") else raw
                        decoded = binascii.unhexlify(hex_raw)
                        params = decoded
                    except Exception:
                        params = raw
        else:
            params = []

    # id is optional (notification when absent)
    id_present = "id" in obj
    req_id = obj.get("id") if id_present else _NO_ID
    if id_present:
        _ = _validate_id(req_id)
    return method, params, req_id


_NO_ID = object()  # sentinel for notification


def _default_ctx() -> Context:
    """Return a minimal Context when callers don't provide one.

    Some integration points historically invoked ``dispatch(payload)`` without
    passing the Context object, which caused ``TypeError: dispatch() missing 1
    required positional argument: 'ctx'``. To remain tolerant of those callers
    (and to keep /rpc from surfacing -32603 Internal error), we synthesize a
    lightweight Context with no request information.
    """

    return Context(request=None, received_at_ms=_now_ms(), client=None, headers={})


async def dispatch_one(obj: Json, ctx: Optional[Context]) -> Optional[Json]:
    ensure_loaded()
    _sync_with_methods_registry()
    """
    Dispatch a single JSON-RPC request object.
    Returns a response object or None (for notifications).
    """
    try:
        method_name, params, req_id = _validate_request_obj(obj)
        alias_map = {
            "tx.send": "tx.sendRawTransaction",
            "tx.broadcast": "tx.sendRawTransaction",
            "tx.submit": "tx.sendRawTransaction",
            "sendRawTransaction": "tx.sendRawTransaction",
            "eth_sendRawTransaction": "tx.sendRawTransaction",
        }
        resolved_method = alias_map.get(method_name, method_name)
        policy = get_active_policy()
        policy.authorize(resolved_method, ctx)
        fn = registry.get(resolved_method)
        args, kwargs = _bind_call_args(fn, params, ctx)
        result = await _invoke_handler(fn, args, kwargs)
        # Notification?
        if req_id is _NO_ID:
            return None
        return {"jsonrpc": "2.0", "id": req_id, "result": result}
    except Exception as exc:
        # Notification? Still no response, even on error per spec (but we *log* it)
        req_id = obj.get("id", _NO_ID)
        method = obj.get("method")
        params = obj.get("params")
        rpc_err = to_error(exc)
        if getattr(rpc_err, "code", None) == -32602:
            existing = dict(getattr(rpc_err, "data", {}) or {})
            existing.setdefault("request", {"id": req_id if req_id is not _NO_ID else None, "method": method})
            existing.setdefault("server_version", getattr(rpc_version, "__version__", "dev"))
            rpc_err = RpcError(code=-32602, message=getattr(rpc_err, "message", "Invalid params"), data=existing)
        if getattr(rpc_err, "code", None) == -32601 and isinstance(method, str):
            suggestions = difflib.get_close_matches(method, registry.names, n=5, cutoff=0.5)
            rpc_err = MethodNotFound(method)
            rpc_err = RpcError(
                code=-32601,
                message="Method not found",
                data={"method": method, "did_you_mean": suggestions},
            )
        if getattr(rpc_err, "code", None) == -32603:
            log.exception(
                "RPC internal error",
                extra={
                    "jsonrpc_method": method,
                    "params": _redact_params(params),
                },
            )
        elif req_id is _NO_ID:
            log.debug("Error in notification %s: %s", method, exc)
            return None
        if req_id is _NO_ID:
            return None
        return {"jsonrpc": "2.0", "id": req_id, "error": _error_obj(rpc_err)}


async def dispatch(
    payload: Union[Json, List[Any]], ctx: Optional[Context] = None
) -> Union[Json, List[Json], None]:
    """
    Dispatch a parsed JSON payload (already json.loads'ed).
    Handles single objects and batches.
    """
    # Backwards-compat: allow callers to omit ctx and fall back to a minimal one
    if ctx is None:
        ctx = _default_ctx()

    if isinstance(payload, list):
        if len(payload) == 0:
            # Empty batch is invalid
            return {
                "jsonrpc": "2.0",
                "id": None,
                "error": _error_obj(InvalidRequest("empty batch")),
            }

        results: List[Optional[Json]] = []
        for obj in payload:
            if isinstance(obj, dict):
                r = await dispatch_one(obj, ctx)
            else:
                r = {
                    "jsonrpc": "2.0",
                    "id": None,
                    "error": _error_obj(InvalidRequest("Request must be an object")),
                }
            results.append(r)

        out = [r for r in results if r is not None]
        return out

    elif isinstance(payload, dict):
        res = await dispatch_one(payload, ctx)
        return res

    else:
        # Entire payload invalid
        return {
            "jsonrpc": "2.0",
            "id": None,
            "error": _error_obj(InvalidRequest("payload must be object or array")),
        }


# --------------------------------------------------------------------------------------
# FastAPI router
# --------------------------------------------------------------------------------------

router = APIRouter()


@router.post("/")
async def jsonrpc_endpoint(request: Request) -> Response:
    """
    HTTP endpoint for JSON-RPC POST.
    """
    try:
        started = time.perf_counter()
        correlation_id = request.headers.get("x-correlation-id") or uuid.uuid4().hex
        body = await request.body()
        try:
            payload = json.loads(body.decode("utf-8"))
        except Exception:
            raise InvalidRequest("malformed JSON")
        payload_preview: Any
        if isinstance(payload, dict):
            payload_preview = {
                "id": payload.get("id"),
                "method": payload.get("method"),
                "params_type": type(payload.get("params")).__name__,
            }
        elif isinstance(payload, list):
            payload_preview = {"batch_size": len(payload)}
        else:
            payload_preview = type(payload).__name__
        peer = _extract_peer_client(request)
        log.info(
            "rpc.request",
            extra={
                "correlation_id": correlation_id,
                "remote": peer[0] if peer else None,
                "user_agent": request.headers.get("user-agent"),
                "preview": payload_preview,
            },
        )
        # Build context
        ctx = Context(
            request=request,
            received_at_ms=_now_ms(),
            client=peer,
            headers={k.lower(): v for k, v in request.headers.items()},
        )
        result = await dispatch(payload, ctx)
        if result is None:
            return Response(status_code=204)
        elapsed_ms = int((time.perf_counter() - started) * 1000)
        log.info("rpc.response", extra={"correlation_id": correlation_id, "elapsed_ms": elapsed_ms})
        return Response(
            content=json.dumps(result, separators=(",", ":")),
            media_type="application/json",
        )
    except JsonRpcError as e:
        # Top-level structural errors
        log.exception("rpc.request_error")
        err = {"jsonrpc": "2.0", "id": None, "error": _error_obj(e)}
        return Response(
            content=json.dumps(err, separators=(",", ":")),
            status_code=400,
            media_type="application/json",
        )
    except Exception as e:  # pragma: no cover
        err = {"jsonrpc": "2.0", "id": None, "error": _error_obj(InternalError(str(e)))}
        return Response(
            content=json.dumps(err, separators=(",", ":")),
            status_code=500,
            media_type="application/json",
        )


# --------------------------------------------------------------------------------------
# Introspection helper (optional)
# --------------------------------------------------------------------------------------


@registry.method("rpc.discover")
async def rpc_discover() -> Dict[str, Any]:
    """Return the OpenRPC document for this server (same as GET /openrpc.json)."""

    doc = _load_openrpc_document()
    try:
        from coretx.crypto import list_scheme_descriptors
        ext = dict(doc.get("extensions") or {})
        ext["signatureSchemes"] = list_scheme_descriptors()
        doc["extensions"] = ext
    except Exception:
        pass
    return doc


@registry.method("rpc.listMethods")
async def rpc_list_methods() -> List[str]:
    """Return the list of registered method names (for debugging/clients)."""
    return registry.names
