"""Framework-agnostic core logic for Notary compliance logging."""

import hashlib
import uuid
from dataclasses import dataclass
from datetime import UTC, datetime
from importlib import metadata
from typing import Any

import httpx
import jcs

from .arweave import ArweaveBackend
from .config import ArweaveHashStorage, CustodiedHashStorage, RawPayloadStorage
from .storage import Storage, create_storage

try:
    __version__ = metadata.version("agentsystems-notary")
except metadata.PackageNotFoundError:
    __version__ = "0.0.0"


# Maximum payload size for Arweave hash storage (10KB)
MAX_PAYLOAD_SIZE = 10 * 1024


class PayloadTooLargeError(Exception):
    """Raised when payload exceeds maximum size limit for Arweave hash storage."""

    pass


@dataclass
class LogResult:
    """Result from logging an interaction."""

    content_hash: str
    custodied_receipt: str | None = None
    arweave_tx_id: str | None = None


# Type alias for hash storage options
HashStorage = CustodiedHashStorage | ArweaveHashStorage


class NotaryCore:
    """
    Framework-agnostic notary logging core.

    Handles canonicalization, hashing, and storage for any AI framework.

    Args:
        raw_payload_storage: Configuration for vendor's S3 bucket (full audit logs)
        hash_storage: List of hash storage configurations (custodied and/or Arweave)
        debug: Enable debug output (default: False)

    Example:
        ```python
        from agentsystems_notary import (
            NotaryCore,
            RawPayloadStorage,
            CustodiedHashStorage,
        )

        raw_payload_storage = RawPayloadStorage(
            bucket_name="my-audit-logs",
            aws_access_key_id="...",
            aws_secret_access_key="...",
        )

        core = NotaryCore(
            raw_payload_storage=raw_payload_storage,
            hash_storage=[
                CustodiedHashStorage(api_key="sk_asn_prod_...", slug="my_tenant"),
            ],
        )

        result = core.log_interaction(
            input_data={"prompt": "Hello"},
            output_data={"response": "Hi there!"},
        )
        print(f"Hash: {result.content_hash}")
        ```
    """

    def __init__(
        self,
        raw_payload_storage: RawPayloadStorage,
        hash_storage: list[HashStorage],
        debug: bool = False,
    ):
        if not hash_storage:
            raise ValueError("At least one hash_storage must be provided")

        self.raw_payload_storage = raw_payload_storage
        self.hash_storage = hash_storage
        self.debug = debug

        # Separate hash storage by type
        self._custodied_storages = [
            h for h in hash_storage if isinstance(h, CustodiedHashStorage)
        ]
        self._arweave_storages = [
            h for h in hash_storage if isinstance(h, ArweaveHashStorage)
        ]

        # Detect test mode from any custodied storage
        self.is_test_mode = any(
            h.api_key.startswith("sk_asn_test_") for h in self._custodied_storages
        )

        # Initialize Arweave backends
        self._arweave_backends = [
            ArweaveBackend(storage, debug) for storage in self._arweave_storages
        ]

        # Initialize storage backend
        self._storage: Storage = create_storage(raw_payload_storage.storage)

        # Session tracking
        self.session_id = str(uuid.uuid4())
        self.sequence = 0

    def log_interaction(
        self,
        input_data: dict[str, Any],
        output_data: dict[str, Any],
        metadata: dict[str, Any] | None = None,
    ) -> LogResult:
        """
        Log an LLM interaction with cryptographic verification.

        This is the main entry point called by framework adapters.
        Performs: canonicalization -> hashing -> storage

        Args:
            input_data: Framework-specific input (prompts, messages, etc.)
            output_data: Framework-specific output (response text, etc.)
            metadata: Additional metadata to include

        Returns:
            LogResult with content hash and receipts/transaction IDs

        Raises:
            PayloadTooLargeError: If payload exceeds 10KB and Arweave is configured
            Exception: If any hash storage fails (strict mode)
        """
        self.sequence += 1

        # Build payload (include slug from first CustodiedHashStorage if present)
        first_custodied = (
            self._custodied_storages[0] if self._custodied_storages else None
        )
        payload = {
            "metadata": {
                "session_id": self.session_id,
                "sequence": self.sequence,
                "timestamp": datetime.now(UTC).isoformat(),
                "slug": first_custodied.slug if first_custodied else None,
                **(metadata or {}),
            },
            "input": input_data,
            "output": output_data,
        }

        # 1. Canonicalize (deterministic JSON serialization)
        canonical_bytes = jcs.canonicalize(payload)

        if self.debug:
            print("\n" + "=" * 80)
            print("DATA_TO_HASH (canonical JSON):")
            print(canonical_bytes.decode("utf-8"))
            print("=" * 80)

        # Check size limit (only if Arweave storage is configured)
        if self._arweave_storages and len(canonical_bytes) > MAX_PAYLOAD_SIZE:
            raise PayloadTooLargeError(
                f"Payload size ({len(canonical_bytes)} bytes) exceeds "
                f"maximum ({MAX_PAYLOAD_SIZE} bytes) for Arweave hash storage"
            )

        # 2. Hash
        content_hash = hashlib.sha256(canonical_bytes).hexdigest()

        if self.debug:
            print(f"HASH: {content_hash}")
            print("=" * 80 + "\n")

        # Initialize result
        result = LogResult(content_hash=content_hash)

        # 3. Write to hash storages (strict mode - any failure raises)
        # Each hash storage gets its own S3 path for clear separation
        for custodied_storage in self._custodied_storages:
            receipt, tenant_id = self._upload_to_custodied(
                custodied_storage, content_hash
            )
            result.custodied_receipt = receipt  # Last one wins if multiple
            # Write full payload to S3 using tenant_id from API
            env = (
                "test"
                if custodied_storage.api_key.startswith("sk_asn_test_")
                else "prod"
            )
            self._write_to_storage(canonical_bytes, content_hash, tenant_id, env)

        for i, arweave_backend in enumerate(self._arweave_backends):
            tx_id = arweave_backend.upload_hash(
                content_hash, self.session_id, self.sequence
            )
            result.arweave_tx_id = tx_id  # Last one wins if multiple
            # Write to S3 using namespace from Arweave storage
            namespace = self._arweave_storages[i].namespace
            self._write_to_storage(canonical_bytes, content_hash, namespace, "arweave")

        return result

    def _upload_to_custodied(
        self,
        storage: CustodiedHashStorage,
        content_hash: str,
    ) -> tuple[str, str]:
        """
        Upload hash to custodied storage API.

        Args:
            storage: CustodiedHashStorage configuration
            content_hash: SHA-256 hash of canonical payload

        Returns:
            Tuple of (receipt, tenant_id)

        Raises:
            httpx.HTTPStatusError: If API request fails
            ValueError: If response missing tenant_id
        """
        with httpx.Client(timeout=5.0) as client:
            resp = client.post(
                storage.api_url,
                headers={
                    "X-API-Key": storage.api_key,
                    "X-SDK-Version": __version__,
                },
                json={"hash": content_hash, "slug": storage.slug},
            )
            resp.raise_for_status()
            result = resp.json()
            receipt = result["receipt"]
            tenant_id = result.get("tenant_id")

        if not tenant_id:
            raise ValueError("Missing tenant_id from custodied storage API response")

        if self.debug:
            print(f"[Custodied] Receipt: {receipt[:8]}...")

        return receipt, tenant_id

    def _write_to_storage(
        self,
        data_bytes: bytes,
        content_hash: str,
        path_prefix: str,
        env_prefix: str,
    ) -> None:
        """
        Write full payload to vendor's storage bucket.

        Path: {env_prefix}/{path_prefix}/{YYYY}/{MM}/{DD}/{hash}.json

        Args:
            data_bytes: Canonical JSON bytes
            content_hash: SHA-256 hash (used in filename)
            path_prefix: tenant_id (custodied) or namespace (Arweave)
            env_prefix: "test", "prod", or "arweave"
        """
        date_path = datetime.now(UTC).strftime("%Y/%m/%d")
        key = f"{env_prefix}/{path_prefix}/{date_path}/{content_hash}.json"

        self._storage.put_object(
            key=key,
            body=data_bytes,
            content_type="application/json",
            metadata={"hash": content_hash},
        )

        if self.debug:
            print(f"[Storage] Saved: {self._storage.bucket_name}/{key}")
