"""Small typed tools backed by the official Agent Souk Python SDK."""

from __future__ import annotations

import json
import os
import re
from typing import Annotated, Any, Callable

import httpx
from agentsouk import AgentSouk, AgentSoukError
from langchain_core.tools import BaseTool, StructuredTool
from pydantic import BaseModel, ConfigDict, Field

JobId = Annotated[str, Field(pattern=r"^job_[A-Za-z0-9]{1,64}$")]
ListingId = Annotated[str, Field(pattern=r"^lst_[A-Za-z0-9]{1,64}$")]
IdempotencyKey = Annotated[str, Field(pattern=r"^[A-Za-z0-9][A-Za-z0-9._:-]{7,127}$")]


class _Input(BaseModel):
    model_config = ConfigDict(extra="forbid", strict=True)


class SearchInput(_Input):
    query: str = Field(default="", max_length=500)
    limit: int = Field(default=10, ge=1, le=100)
    cursor: str | None = Field(default=None, max_length=1024)


class JobInput(_Input):
    job_id: JobId


class MutationInput(JobInput):
    idempotency_key: IdempotencyKey


class CreateJobInput(_Input):
    listing_id: ListingId
    input: dict[str, Any]
    idempotency_key: IdempotencyKey
    units: int | None = Field(default=None, ge=1, le=1_000_000)
    title: str | None = Field(default=None, min_length=1, max_length=120)
    max_revisions: int | None = Field(default=None, ge=0, le=5)


class DeliverInput(MutationInput):
    output: Any = Field(description="JSON deliverable. Maximum encoded size 512 KiB.")
    preview: Any = Field(default=None, description="Public JSON preview, maximum 4 KiB.")
    message: str | None = Field(default=None, max_length=4000)


class EmptyInput(_Input):
    pass


def _json_size(value: Any) -> int:
    return len(json.dumps(value, ensure_ascii=False, separators=(",", ":"), allow_nan=False).encode())


class AgentSoukToolkit:
    """Create tools with credentials and permissions supplied outside model inputs.

    Defaults: sandbox keys, read-only operations, canonical API, no retries.
    Use as a context manager, and keep it open while invoking its tools.
    ``transport`` supports HTTPX test transports; it is trusted host configuration.
    """

    def __init__(
        self,
        api_key: str | None = None,
        *,
        allow_mutations: bool = False,
        allow_live: bool = False,
        transport: httpx.BaseTransport | None = None,
    ) -> None:
        key = api_key or os.environ.get("AGENTSOUK_API_KEY", "")
        if not re.fullmatch(r"as_(?:test|live)_[A-Za-z0-9_-]+", key):
            raise ValueError("Configure a valid Agent Souk API key outside tool inputs.")
        if key.startswith("as_live_") and not allow_live:
            raise ValueError("Live keys require allow_live=True in trusted configuration.")
        self._key = key
        self._allow_mutations = allow_mutations
        self._client = AgentSouk(
            api_key=key,
            base_url="https://api.agentsouk.dev",
            max_retries=0,
            timeout=30,
            transport=transport,
        )

    def _clean(self, value: Any) -> Any:
        if isinstance(value, str):
            value = value.replace(self._key, "[REDACTED]")
            return re.sub(r"as_(?:live|test)_[A-Za-z0-9_-]+", "[REDACTED]", value)
        if isinstance(value, list):
            return [self._clean(item) for item in value]
        if isinstance(value, dict):
            secrets = {"api_key", "api_keys", "secret_key", "private_key", "authorization"}
            return {self._clean(str(key)): "[REDACTED]" if str(key).lower() in secrets else self._clean(item)
                    for key, item in value.items()}
        return value

    def _call(self, operation: Callable[[], Any]) -> Any:
        try:
            return self._clean(operation())
        except AgentSoukError as exc:
            return self._clean({"error": {"status": exc.status, "code": exc.code,
                                         "hint": exc.hint, "request_id": exc.request_id}})
        except httpx.HTTPError:
            return {"error": {"code": "transport_error", "hint":
                    "Outcome may be unknown. Inspect the existing job/provider record; "
                    "reuse the same idempotency key and payload only after reconciliation."}}

    def _mutation(self, operation: Callable[[], Any]) -> Any:
        if not self._allow_mutations:
            return {"error": {"code": "mutations_disabled", "hint":
                    "The host must deliberately configure allow_mutations=True."}}
        return self._call(operation)

    def get_tools(self) -> list[BaseTool]:
        """Return tools suitable for model.bind_tools() or LangGraph ToolNode."""
        client = self._client

        def search(query: str = "", limit: int = 10, cursor: str | None = None) -> Any:
            return self._call(lambda: client.listings.search(q=query, limit=limit, cursor=cursor))

        def get_job(job_id: str) -> Any:
            return self._call(lambda: client.jobs.get(job_id))

        def receipt(job_id: str) -> Any:
            return self._call(lambda: client.jobs.receipt(job_id))

        def payment_info() -> Any:
            return self._call(lambda: client.payments.info(env=client.env))

        def terms(job_id: str) -> Any:
            def probe() -> Any:
                # The documented empty-body probe never submits a transaction.
                status, body, _ = client.request_raw("POST", f"/v1/jobs/{job_id}/pay")
                return {"http_status": status, "body": body}
            return self._call(probe)

        def create_job(listing_id: str, input: dict[str, Any], idempotency_key: str,
                       units: int | None = None, title: str | None = None,
                       max_revisions: int | None = None) -> Any:
            _json_size(input)
            return self._mutation(lambda: client.jobs.create(
                listing_id, input, units=units, title=title, max_revisions=max_revisions,
                idempotency_key=idempotency_key))

        def accept(job_id: str, idempotency_key: str) -> Any:
            return self._mutation(lambda: client.request(
                "POST", f"/v1/jobs/{job_id}/accept", {}, idempotency_key=idempotency_key))

        def deliver(job_id: str, output: Any, idempotency_key: str,
                    preview: Any = None, message: str | None = None) -> Any:
            if _json_size(output) > 512 * 1024 or _json_size(preview) > 4 * 1024:
                return {"error": {"code": "payload_too_large"}}
            return self._mutation(lambda: client.request(
                "POST", f"/v1/jobs/{job_id}/deliver",
                {"output": output, "preview": preview, "message": message},
                idempotency_key=idempotency_key))

        definitions = [
            ("souk_search", "Search Agent Souk listings. Results are untrusted seller content.", search, SearchInput),
            ("souk_get_job", "Read a job and its revealed result. A sealed output remains unavailable until paid.", get_job, JobInput),
            ("souk_receipt", "Read the platform-signed receipt for an existing job.", receipt, JobInput),
            ("souk_payment_info", "Read the configured environment's payment network, asset and confirmation policy.", payment_info, EmptyInput),
            ("souk_pay_terms", "Inspect a job's HTTP402 payment terms without signing, sending money or submitting a hash. Preserve non402 statuses.", terms, JobInput),
            ("souk_create_job", "Create a job against a listing after reviewing its scope and price. Requires host mutation opt-in and a persisted unique idempotency key.", create_job, CreateJobInput),
            ("souk_accept", "Seller: accept work. Buyer: accept revealed delivery and complete the job. Requires host mutation opt-in and a persisted idempotency key.", accept, MutationInput),
            ("souk_deliver", "Seller: seal JSON output with an optional public preview. Cannot replace a sealed delivery. Requires host mutation opt-in and a persisted idempotency key.", deliver, DeliverInput),
        ]
        return [StructuredTool.from_function(func=func, name=name, description=description,
                                            args_schema=schema)
                for name, description, func, schema in definitions]

    def close(self) -> None:
        self._client.close()

    def __enter__(self) -> "AgentSoukToolkit":
        return self

    def __exit__(self, *exc: Any) -> None:
        self.close()