"""ProofTrail SDK: decorator-based tamper-evident audit trails for tool calls.

Provides a simple API for creating cryptographic receipts for every function
call an AI agent makes. Built on ads_foundation's Merkle tree primitives.

Usage:
    from ads_foundation.prooftrail import ProofTrailMonitor

    monitor = ProofTrailMonitor()

    @monitor.track_tool_call
    def query_db(customer_id: str) -> dict:
        return {"balance": 100.0}

    with monitor.session(agent_id="support_bot") as session:
        result = query_db("cust_123")
        receipt = session.get_latest_receipt()
        assert session.verify_receipt(receipt)
"""

from __future__ import annotations

import functools
import json
import logging
import queue
import threading
from contextvars import ContextVar
from dataclasses import dataclass
from datetime import datetime, timezone

from .hash import sha256
from .merkle import (
    MerkleProof,
    Side,
    build_merkle,
    generate_proof,
    root_hash,
    verify_proof as merkle_verify,
)

logger = logging.getLogger(__name__)

DEFAULT_API_URL = "https://prooftrail.up.railway.app"

# Context variable for the active session (supports async / nested sessions).
_active_session: ContextVar[Session | None] = ContextVar(
    "prooftrail_session", default=None
)

GENESIS_HASH: bytes = b"\x00" * 32


def _serialize(obj: object) -> bytes:
    """Deterministic JSON serialization for hashing."""
    return json.dumps(obj, sort_keys=True, default=str).encode()


@dataclass(frozen=True)
class Receipt:
    """Cryptographic receipt for a single tracked tool call."""

    index: int
    timestamp: str
    agent_id: str | None
    tool_name: str
    query_hash: bytes
    response_hash: bytes
    record_hash: bytes
    root_hash: bytes
    proof: MerkleProof


class Session:
    """Scoped session binding an agent_id to tracked calls.

    Use as a context manager via ``monitor.session(agent_id=...)``.
    """

    __slots__ = ("agent_id", "_monitor", "_receipts", "_token")

    def __init__(self, agent_id: str, monitor: ProofTrailMonitor) -> None:
        self.agent_id = agent_id
        self._monitor = monitor
        self._receipts: list[Receipt] = []
        self._token = None

    # -- context manager --------------------------------------------------

    def __enter__(self) -> Session:
        self._token = _active_session.set(self)
        return self

    def __exit__(self, *exc) -> None:
        _active_session.reset(self._token)
        return None

    # -- internal ---------------------------------------------------------

    def _record(self, receipt: Receipt) -> None:
        self._receipts.append(receipt)

    # -- public API -------------------------------------------------------

    def get_latest_receipt(self) -> Receipt:
        """Return the most recent receipt in this session."""
        if not self._receipts:
            raise ValueError("No receipts in this session")
        return self._receipts[-1]

    def get_receipt(self, index: int) -> Receipt:
        """Return the receipt at the given chain index."""
        for r in self._receipts:
            if r.index == index:
                return r
        raise KeyError(f"No receipt with index {index} in this session")

    def verify_receipt(self, receipt: Receipt) -> bool:
        """Verify a receipt's Merkle proof against its root hash."""
        return self._monitor.verify_receipt(receipt)

    def trail(self) -> list[Receipt]:
        """Return all receipts in this session, in order."""
        return list(self._receipts)


def _receipt_to_submission(receipt: Receipt) -> dict:
    """Convert a Receipt to the dict format expected by POST /receipts."""
    return {
        "index": receipt.index,
        "timestamp": receipt.timestamp,
        "agent_id": receipt.agent_id,
        "tool_name": receipt.tool_name,
        "query_hash": receipt.query_hash.hex(),
        "response_hash": receipt.response_hash.hex(),
        "record_hash": receipt.record_hash.hex(),
        "root_hash": receipt.root_hash.hex(),
        "proof": [
            {"side": side.name, "sibling_hash": sibling.hex()}
            for side, sibling in receipt.proof.path
        ],
        "proof_root": receipt.proof.root_hash.hex(),
    }


class _CloudSubmitter:
    """Background daemon thread that submits receipts to the hosted API."""

    _MAX_QUEUE = 1000

    def __init__(self, api_key: str, api_url: str) -> None:
        self._api_key = api_key
        self._api_url = api_url.rstrip("/")
        self._queue: queue.Queue[dict | None] = queue.Queue(maxsize=self._MAX_QUEUE)
        self._thread = threading.Thread(target=self._worker, daemon=True)
        self._thread.start()

    def submit(self, payload: dict) -> None:
        try:
            self._queue.put_nowait(payload)
        except queue.Full:
            logger.warning("ProofTrail cloud queue full, dropping receipt")

    def _worker(self) -> None:
        import httpx

        client = httpx.Client(
            base_url=self._api_url,
            headers={"X-API-Key": self._api_key},
            timeout=10.0,
        )
        try:
            while True:
                payload = self._queue.get()
                if payload is None:
                    break
                try:
                    resp = client.post("/receipts", json=payload)
                    if resp.status_code >= 400:
                        logger.warning(
                            "ProofTrail cloud submission failed: %d %s",
                            resp.status_code,
                            resp.text,
                        )
                except Exception:
                    logger.warning("ProofTrail cloud submission error", exc_info=True)
        finally:
            client.close()

    def shutdown(self) -> None:
        self._queue.put(None)
        self._thread.join(timeout=5.0)


class ProofTrailMonitor:
    """Main entry point for ProofTrail.

    When ``api_key`` is provided and ``httpx`` is installed, receipts are
    submitted to the hosted API in the background (fire-and-forget).
    """

    __slots__ = ("_api_key", "_api_url", "_chain", "_record_hashes", "_submitter")

    def __init__(
        self,
        api_key: str | None = None,
        api_url: str | None = None,
    ) -> None:
        self._api_key = api_key
        self._api_url = api_url or DEFAULT_API_URL
        self._chain: list[bytes] = []  # ordered record hashes
        self._record_hashes: list[bytes] = []  # same, for Merkle tree
        self._submitter: _CloudSubmitter | None = None

        if api_key:
            try:
                import httpx  # noqa: F401

                self._submitter = _CloudSubmitter(api_key, self._api_url)
            except ImportError:
                logger.warning(
                    "httpx not installed -- receipts stored locally only. "
                    "Install with: pip install ads-foundation[cloud]"
                )

    # -- decorator --------------------------------------------------------

    def track_tool_call(self, fn):
        """Decorator that records every call made inside an active session."""

        @functools.wraps(fn)
        def wrapper(*args, **kwargs):
            session = _active_session.get(None)
            if session is None or session._monitor is not self:
                return fn(*args, **kwargs)

            # Serialize inputs
            call_data = {"args": args, "kwargs": kwargs}
            query_hash = sha256(_serialize(call_data))

            # Execute the real function
            result = fn(*args, **kwargs)

            # Serialize output
            response_hash = sha256(_serialize(result))

            # Build chained record hash
            timestamp = datetime.now(timezone.utc).isoformat()
            prev_hash = self._chain[-1] if self._chain else GENESIS_HASH
            record_hash = sha256(
                timestamp.encode()
                + (session.agent_id or "").encode()
                + fn.__name__.encode()
                + query_hash
                + response_hash
                + prev_hash
            )

            # Append to chain
            self._chain.append(record_hash)
            self._record_hashes.append(record_hash)

            # Build Merkle tree and generate proof
            tree = build_merkle(self._record_hashes)
            index = len(self._record_hashes) - 1
            proof = generate_proof(index, tree, len(self._record_hashes))

            receipt = Receipt(
                index=index,
                timestamp=timestamp,
                agent_id=session.agent_id,
                tool_name=fn.__name__,
                query_hash=query_hash,
                response_hash=response_hash,
                record_hash=record_hash,
                root_hash=root_hash(tree),
                proof=proof,
            )
            session._record(receipt)

            if self._submitter:
                self._submitter.submit(_receipt_to_submission(receipt))

            return result

        return wrapper

    # -- session ----------------------------------------------------------

    def session(self, agent_id: str) -> Session:
        """Create a new session context for the given agent."""
        return Session(agent_id=agent_id, monitor=self)

    # -- verification -----------------------------------------------------

    def verify_receipt(self, receipt: Receipt) -> bool:
        """Verify a receipt's Merkle proof against its embedded root hash."""
        return merkle_verify(receipt.record_hash, receipt.proof)
