"""Clampd Python SDK — Guard AI agent tool calls in 1 line.

Usage:
    import clampd

    # Guard OpenAI tool calls
    client = clampd.openai(OpenAI(), agent_id="my-agent")

    # Guard Anthropic/Claude tool calls
    client = clampd.anthropic(Anthropic(), agent_id="my-agent")

    # Guard any function
    @clampd.guard("database.query", agent_id="my-agent")
    def run_query(sql: str) -> str: ...

    # Guard LangChain agents (attach as callback)
    agent.invoke(input, config={"callbacks": [clampd.langchain(agent_id="my-agent")]})

    # Guard Google ADK agents
    agent = Agent(tools=[...], before_tool_callback=clampd.adk(agent_id="my-agent"))
"""

from __future__ import annotations

import functools
import inspect
import json
import logging
import os
from typing import Any, Callable, TypeVar

from clampd.auth import make_agent_jwt
from clampd.client import ClampdClient, ClampdBlockedError, ProxyResponse

__all__ = [
    "ClampdClient",
    "ClampdBlockedError",
    "ProxyResponse",
    "make_agent_jwt",
    "guard",
    "openai",
    "anthropic",
    "langchain",
    "adk",
    "init",
]

F = TypeVar("F", bound=Callable[..., Any])

logger = logging.getLogger("clampd")

# ── Global config ─────────────────────────────────────────────────────

_default_client: ClampdClient | None = None


def init(
    *,
    agent_id: str,
    gateway_url: str = "http://localhost:8080",
    api_key: str = "clmpd_demo_key",
    secret: str | None = None,
) -> ClampdClient:
    """Initialize the global Clampd client. Call once at startup.

    After calling ``init()``, other functions like ``guard()``,
    ``openai()``, etc. will use this client automatically.

        clampd.init(agent_id="my-agent")
    """
    global _default_client
    _default_client = ClampdClient(
        gateway_url=gateway_url,
        agent_id=agent_id,
        api_key=api_key,
        secret=secret,
    )
    return _default_client


def _get_client(
    agent_id: str | None = None,
    gateway_url: str | None = None,
    api_key: str | None = None,
    secret: str | None = None,
) -> ClampdClient:
    """Get or create a ClampdClient."""
    if _default_client is not None:
        return _default_client

    if not agent_id:
        agent_id = os.environ.get("CLAMPD_AGENT_ID", "")
        if not agent_id:
            raise ValueError(
                "No agent_id provided. Either call clampd.init(agent_id=...) "
                "first, or pass agent_id= to each function, or set CLAMPD_AGENT_ID env var."
            )

    return ClampdClient(
        gateway_url=gateway_url or os.environ.get("CLAMPD_GATEWAY_URL", "http://localhost:8080"),
        agent_id=agent_id,
        api_key=api_key or os.environ.get("CLAMPD_API_KEY", "clmpd_demo_key"),
        secret=secret,
    )


# ── Response inspection helper ────────────────────────────────────────


def _inspect_response(
    client: ClampdClient,
    tool_name: str,
    response_data: Any,
    request_id: str = "",
    fail_open: bool = False,
) -> None:
    """Send a tool response through the inspect endpoint."""
    try:
        serializable = response_data if isinstance(response_data, (str, int, float, bool, type(None), dict, list)) else str(response_data)
        resp = client.inspect(tool=tool_name, response_data=serializable, request_id=request_id)
    except Exception as e:
        if fail_open:
            logger.warning("Clampd inspect error (fail-open): %s", e)
            return
        raise ClampdBlockedError(f"Response inspection failed: {e}") from e

    if not resp.allowed:
        raise ClampdBlockedError(
            resp.denial_reason or "Response blocked",
            risk_score=resp.risk_score,
            response=resp,
        )


# ── @clampd.guard() decorator ────────────────────────────────────────


def guard(
    tool_name: str,
    *,
    agent_id: str | None = None,
    target_url: str = "",
    fail_open: bool = False,
    check_response: bool = False,
) -> Callable[[F], F]:
    """Decorator that guards any function through the Clampd pipeline.

        @clampd.guard("database.query")
        def run_query(sql: str) -> str:
            return db.execute(sql)

    Set ``check_response=True`` to also inspect the return value for PII,
    data anomalies, or policy violations.
    """
    client = _get_client(agent_id=agent_id)

    def decorator(fn: F) -> F:
        sig = inspect.signature(fn)

        @functools.wraps(fn)
        def wrapper(*args: Any, **kwargs: Any) -> Any:
            bound = sig.bind(*args, **kwargs)
            bound.apply_defaults()
            params = dict(bound.arguments)

            try:
                resp = client.proxy(tool=tool_name, params=params, target_url=target_url)
            except Exception as e:
                if fail_open:
                    logger.warning("Clampd gateway error (fail-open): %s", e)
                    return fn(*args, **kwargs)
                raise ClampdBlockedError(str(e)) from e

            if not resp.allowed:
                raise ClampdBlockedError(
                    resp.denial_reason or "denied",
                    risk_score=resp.risk_score,
                    response=resp,
                )

            result = fn(*args, **kwargs)

            if check_response:
                _inspect_response(client, tool_name, result, resp.request_id, fail_open)

            return result

        return wrapper  # type: ignore[return-value]

    return decorator


# ── clampd.openai() — wrap OpenAI client ──────────────────────────────


def openai(
    client: Any,
    *,
    agent_id: str | None = None,
    target_url: str = "",
    fail_open: bool = False,
    check_response: bool = False,
) -> Any:
    """Wrap an OpenAI client so all tool calls go through Clampd.

        import openai, clampd
        client = clampd.openai(openai.OpenAI(), agent_id="my-agent")
        # Use client.chat.completions.create() as normal — tool calls are guarded

    Set ``check_response=True`` to also inspect tool responses for PII or anomalies.
    Returns a drop-in replacement that intercepts tool execution.
    """
    clampd_client = _get_client(agent_id=agent_id)
    original_create = client.chat.completions.create

    def guarded_create(*args: Any, **kwargs: Any) -> Any:
        response = original_create(*args, **kwargs)
        choice = response.choices[0]

        if choice.finish_reason != "tool_calls" or not choice.message.tool_calls:
            return response

        for tc in choice.message.tool_calls:
            tool_args = json.loads(tc.function.arguments) if isinstance(tc.function.arguments, str) else tc.function.arguments
            try:
                result = clampd_client.proxy(
                    tool=tc.function.name,
                    params=tool_args,
                    target_url=target_url,
                )
            except Exception as e:
                if fail_open:
                    continue
                raise ClampdBlockedError(str(e)) from e

            if not result.allowed:
                raise ClampdBlockedError(
                    result.denial_reason or "denied",
                    risk_score=result.risk_score,
                    response=result,
                )

        return response

    client.chat.completions.create = guarded_create

    if check_response:
        client._clampd_inspect = lambda tool, data, req_id="": _inspect_response(
            clampd_client, tool, data, req_id, fail_open
        )

    return client


# ── clampd.anthropic() — wrap Anthropic client ───────────────────────


def anthropic(
    client: Any,
    *,
    agent_id: str | None = None,
    target_url: str = "",
    fail_open: bool = False,
    check_response: bool = False,
) -> Any:
    """Wrap an Anthropic client so all tool calls go through Clampd.

        import anthropic, clampd
        client = clampd.anthropic(anthropic.Anthropic(), agent_id="my-agent")
        # Use client.messages.create() as normal — tool_use blocks are guarded

    Set ``check_response=True`` to also inspect tool responses for PII or anomalies.
    Returns a drop-in replacement that intercepts tool_use blocks.
    """
    clampd_client = _get_client(agent_id=agent_id)
    original_create = client.messages.create

    def guarded_create(*args: Any, **kwargs: Any) -> Any:
        response = original_create(*args, **kwargs)

        if response.stop_reason != "tool_use":
            return response

        for block in response.content:
            if block.type != "tool_use":
                continue

            tool_args = block.input if isinstance(block.input, dict) else {}
            try:
                result = clampd_client.proxy(
                    tool=block.name,
                    params=tool_args,
                    target_url=target_url,
                )
            except Exception as e:
                if fail_open:
                    continue
                raise ClampdBlockedError(str(e)) from e

            if not result.allowed:
                raise ClampdBlockedError(
                    result.denial_reason or "denied",
                    risk_score=result.risk_score,
                    response=result,
                )

        return response

    client.messages.create = guarded_create

    if check_response:
        client._clampd_inspect = lambda tool, data, req_id="": _inspect_response(
            clampd_client, tool, data, req_id, fail_open
        )

    return client


# ── clampd.langchain() — callback handler ────────────────────────────


def langchain(
    *,
    agent_id: str | None = None,
    target_url: str = "",
    fail_open: bool = False,
) -> Any:
    """Create a LangChain callback handler that guards all tool calls.

        agent.invoke(input, config={"callbacks": [clampd.langchain(agent_id="my-agent")]})

    Or attach globally:

        from langchain_core.globals import set_llm_cache
        callbacks = [clampd.langchain(agent_id="my-agent")]
    """
    from clampd.langchain_callback import ClampdCallbackHandler

    client = _get_client(agent_id=agent_id)
    return ClampdCallbackHandler(client, target_url=target_url, fail_open=fail_open)


# ── clampd.adk() — Google ADK before_tool_callback ───────────────────


def adk(
    *,
    agent_id: str | None = None,
    target_url: str = "",
    fail_open: bool = False,
    check_response: bool = False,
) -> Callable | tuple[Callable, Callable]:
    """Create Google ADK before_tool_callback (and optionally after_tool_callback).

        agent = Agent(
            model="gemini-2.0-flash",
            tools=[search, calculator],
            before_tool_callback=clampd.adk(agent_id="my-agent"),
        )

    With ``check_response=True``, returns a tuple of (before_tool, after_tool):

        before_cb, after_cb = clampd.adk(agent_id="my-agent", check_response=True)
        agent = Agent(..., before_tool_callback=before_cb, after_tool_callback=after_cb)

    Returns None to allow, or a response dict to block.
    """
    client = _get_client(agent_id=agent_id)

    def before_tool(tool_name: str, args: dict[str, Any], context: Any) -> dict[str, Any] | None:
        try:
            result = client.proxy(tool=tool_name, params=args, target_url=target_url)
        except Exception as e:
            if fail_open:
                return None
            return {"error": f"Security gateway error: {e}"}

        if not result.allowed:
            return {"error": f"Blocked by Clampd: {result.denial_reason}", "risk_score": result.risk_score}

        return None

    def after_tool(tool_name: str, response: Any, context: Any) -> dict[str, Any] | None:
        try:
            serializable = response if isinstance(response, (str, int, float, bool, type(None), dict, list)) else str(response)
            result = client.inspect(tool=tool_name, response_data=serializable)
        except Exception as e:
            if fail_open:
                return None
            return {"error": f"Response inspection error: {e}"}

        if not result.allowed:
            return {"error": f"Response blocked by Clampd: {result.denial_reason}", "risk_score": result.risk_score}

        return None

    if check_response:
        return before_tool, after_tool
    return before_tool
