from __future__ import annotations

import hashlib
import json
import os
from dataclasses import dataclass
from typing import Any, Optional, Dict

import httpx


def _get_api_url() -> str:
    return os.getenv("TREESHIP_API_URL", "https://api.treeship.dev")


@dataclass
class AttestResult:
    """Result of creating an attestation."""
    
    attestation_id: str
    signature: str
    payload_hash: str
    key_id: str
    timestamp: str
    url: str
    verify_command: str
    agent_slug: str
    action: str
    inputs_hash: str
    is_public: bool = True
    has_zk_proof: bool = False
    # Verification status
    agent_authenticated: Optional[bool] = None
    code_hash_verified: Optional[bool] = None
    tools_authorized: Optional[bool] = None
    human_approved: bool = False

    @classmethod
    def from_dict(cls, d: Dict[str, Any]) -> AttestResult:
        return cls(
            attestation_id=d["attestation_id"],
            signature=d["signature"],
            payload_hash=d["payload_hash"],
            key_id=d["key_id"],
            timestamp=d["timestamp"],
            url=d["public_url"],
            verify_command=d["verify_command"],
            agent_slug=d["agent_slug"],
            action=d["action"],
            inputs_hash=d["inputs_hash"],
            is_public=d.get("is_public", True),
            has_zk_proof=d.get("has_zk_proof", False),
            agent_authenticated=d.get("agent_authenticated"),
            code_hash_verified=d.get("code_hash_verified"),
            tools_authorized=d.get("tools_authorized"),
            human_approved=d.get("human_approved", False),
        )


@dataclass
class AgentIdentity:
    """Agent identity information."""
    
    slug: str
    name: Optional[str]
    identity_score: int
    verified_proofs: list
    code_hash: Optional[str]
    model_id: Optional[str]
    verified_domain: Optional[str]
    tool_manifest_hash: Optional[str]
    tools: Optional[list]
    
    @classmethod
    def from_dict(cls, d: Dict[str, Any]) -> "AgentIdentity":
        return cls(
            slug=d["slug"],
            name=d.get("name"),
            identity_score=d["identity_score"],
            verified_proofs=d.get("verified_proofs", []),
            code_hash=d.get("code_hash"),
            model_id=d.get("model_id"),
            verified_domain=d.get("verified_domain"),
            tool_manifest_hash=d.get("tool_manifest_hash"),
            tools=d.get("tools"),
        )


class Treeship:
    """
    Treeship client for creating cryptographic attestations.
    
    Example:
        >>> from treeship import Treeship
        >>> ts = Treeship()
        >>> result = ts.attest(
        ...     action="Approved loan application #12345",
        ...     inputs_hash=ts.hash({"customer_id": "cust_123"})
        ... )
        >>> print(result.url)
    """

    def __init__(
        self,
        api_key: Optional[str] = None,
        agent: Optional[str] = None,
        api_url: Optional[str] = None,
    ):
        """
        Initialize the Treeship client.
        
        Args:
            api_key: API key for authentication. Defaults to TREESHIP_API_KEY env var.
            agent: Default agent slug. Defaults to TREESHIP_AGENT env var.
            api_url: API base URL. Defaults to TREESHIP_API_URL env var or https://api.treeship.dev.
        """
        self.api_key = api_key or os.environ.get("TREESHIP_API_KEY", "")
        self.agent = agent or os.getenv("TREESHIP_AGENT", "")
        self._api_url = api_url or _get_api_url()
        self._client = httpx.Client(
            base_url=self._api_url,
            headers={
                "Authorization": f"Bearer {self.api_key}",
                "User-Agent": "treeship-sdk/0.1.0",
            },
            timeout=30.0,
        )

    @staticmethod
    def hash(data: Any) -> str:
        """
        Create a SHA256 hash of data for use as inputs_hash.
        
        Args:
            data: Any JSON-serializable data, string, or bytes.
            
        Returns:
            Hex-encoded SHA256 hash.
        """
        if isinstance(data, (dict, list)):
            data = json.dumps(data, sort_keys=True, separators=(",", ":"))
        if isinstance(data, str):
            data = data.encode()
        return hashlib.sha256(data).hexdigest()

    def attest(
        self,
        action: str,
        inputs_hash: Optional[str] = None,
        agent: Optional[str] = None,
        metadata: Optional[Dict[str, Any]] = None,
        is_public: Optional[bool] = None,
        enable_zk: Optional[bool] = None,
        # Authenticity fields
        agent_signature: Optional[str] = None,
        code_hash: Optional[str] = None,
        model_id: Optional[str] = None,
        # Tool verification
        tools_used: Optional[list] = None,
        # Human approval
        human_approved: Optional[bool] = None,
        human_approver_hash: Optional[str] = None,
    ) -> AttestResult:
        """
        Create a cryptographic attestation with optional verification.
        
        Args:
            action: Human-readable description of the action.
            inputs_hash: SHA256 hash of inputs. If not provided, hashes the action string.
            agent: Agent slug. Uses default agent if not provided.
            metadata: Optional key-value metadata.
            is_public: Override org default visibility. True=public, False=private.
            enable_zk: Enable ZK proof for privacy-preserving public verification.
            agent_signature: Agent's Ed25519 signature (proves authenticity).
            code_hash: Current code hash (verified against registered hash).
            model_id: Model version used (e.g., "gpt-4-turbo").
            tools_used: List of tools used (verified against tool manifest).
            human_approved: Whether a human approved this action.
            human_approver_hash: Hash of human approver identity.
            
        Returns:
            AttestResult with attestation details, verification URL, and verification status.
        """
        slug = agent or self.agent
        if not slug:
            raise ValueError(
                "Agent slug required: pass agent= or set TREESHIP_AGENT environment variable"
            )

        if not self.api_key:
            raise ValueError(
                "API key required: pass api_key= or set TREESHIP_API_KEY environment variable"
            )

        payload: Dict[str, Any] = {
            "agent_slug": slug,
            "action": action,
            "inputs_hash": inputs_hash or self.hash(action),
            "metadata": metadata or {},
        }
        
        if is_public is not None:
            payload["is_public"] = is_public
        if enable_zk is not None:
            payload["enable_zk"] = enable_zk
        if agent_signature is not None:
            payload["agent_signature"] = agent_signature
        if code_hash is not None:
            payload["code_hash"] = code_hash
        if model_id is not None:
            payload["model_id"] = model_id
        if tools_used is not None:
            payload["tools_used"] = tools_used
        if human_approved is not None:
            payload["human_approved"] = human_approved
        if human_approver_hash is not None:
            payload["human_approver_hash"] = human_approver_hash

        response = self._client.post("/v1/attest", json=payload)
        response.raise_for_status()
        return AttestResult.from_dict(response.json())

    def verify(self, attestation_id: str) -> Dict[str, Any]:
        """
        Verify an attestation by ID.
        
        Args:
            attestation_id: The attestation UUID.
            
        Returns:
            Dict with verification result including 'valid' boolean.
            
        Raises:
            httpx.HTTPStatusError: If the attestation is not found or request fails.
        """
        response = self._client.get(f"/v1/verify/{attestation_id}")
        response.raise_for_status()
        return response.json()

    def get_agent(self, slug: Optional[str] = None) -> Dict[str, Any]:
        """
        Get agent feed with recent attestations.
        
        Args:
            slug: Agent slug. Uses default agent if not provided.
            
        Returns:
            Dict with agent info and attestations list.
        """
        agent_slug = slug or self.agent
        if not agent_slug:
            raise ValueError("Agent slug required")
        response = self._client.get(f"/v1/agent/{agent_slug}")
        response.raise_for_status()
        return response.json()

    def get_pubkey(self) -> Dict[str, Any]:
        """
        Get the public signing key for independent verification.
        
        Returns:
            Dict with key_id, algorithm, and public_key_pem.
        """
        response = self._client.get("/v1/pubkey")
        response.raise_for_status()
        return response.json()

    def update_settings(
        self,
        default_public: Optional[bool] = None,
        default_zk_enabled: Optional[bool] = None,
        org_name: Optional[str] = None,
    ) -> Dict[str, Any]:
        """
        Update organization-level settings for your API key.
        
        Args:
            default_public: If True, new attestations are public by default.
                           If False, new attestations are private (owner-only) by default.
            default_zk_enabled: If True, enable ZK proofs by default.
            org_name: Organization name for display.
            
        Returns:
            Dict with updated settings.
        """
        payload: Dict[str, Any] = {}
        if default_public is not None:
            payload["default_public"] = default_public
        if default_zk_enabled is not None:
            payload["default_zk_enabled"] = default_zk_enabled
        if org_name is not None:
            payload["org_name"] = org_name
            
        response = self._client.patch(
            "/v1/keys/settings",
            params={"api_key": self.api_key},
            json=payload,
        )
        response.raise_for_status()
        return response.json()

    # ========================================================================
    # Identity Management
    # ========================================================================

    def register_identity(
        self,
        agent: Optional[str] = None,
        name: Optional[str] = None,
        description: Optional[str] = None,
        public_key_pem: Optional[str] = None,
        code_hash: Optional[str] = None,
        model_id: Optional[str] = None,
        domain: Optional[str] = None,
    ) -> Dict[str, Any]:
        """
        Register agent identity with cryptographic proofs.
        
        Args:
            agent: Agent slug. Uses default agent if not provided.
            name: Human-readable agent name.
            description: Agent description.
            public_key_pem: Ed25519 public key in PEM format for self-signing.
            code_hash: SHA256 hash of agent code/model weights.
            model_id: AI model identifier (e.g., "gpt-4-turbo", "claude-3-opus").
            domain: Domain to verify ownership of.
            
        Returns:
            Registration result with identity score and verification challenges.
        """
        slug = agent or self.agent
        if not slug:
            raise ValueError("Agent slug required")
        
        payload = {"slug": slug}
        if name:
            payload["name"] = name
        if description:
            payload["description"] = description
        if public_key_pem:
            payload["public_key_pem"] = public_key_pem
        if code_hash:
            payload["code_hash"] = code_hash
        if model_id:
            payload["model_id"] = model_id
        if domain:
            payload["domain"] = domain
        
        response = self._client.post("/v1/identity/register", json=payload)
        response.raise_for_status()
        return response.json()

    def get_identity(self, agent: Optional[str] = None) -> AgentIdentity:
        """
        Get agent identity information.
        
        Args:
            agent: Agent slug. Uses default agent if not provided.
            
        Returns:
            AgentIdentity with verification status and proofs.
        """
        slug = agent or self.agent
        if not slug:
            raise ValueError("Agent slug required")
        
        response = self._client.get(f"/v1/identity/{slug}")
        response.raise_for_status()
        return AgentIdentity.from_dict(response.json())

    def set_tool_manifest(
        self,
        agent: Optional[str] = None,
        tools: Optional[list] = None,
        tool_descriptions: Optional[Dict[str, str]] = None,
    ) -> Dict[str, Any]:
        """
        Set authorized tools for an agent.
        
        Example:
            >>> ts.set_tool_manifest(
            ...     agent="my-agent",
            ...     tools=["read_file", "write_file", "send_email"],
            ...     tool_descriptions={"send_email": "Approved domains only"}
            ... )
        
        Args:
            agent: Agent slug. Uses default agent if not provided.
            tools: List of authorized tool names.
            tool_descriptions: Optional descriptions for each tool.
            
        Returns:
            Tool manifest with hash and signature.
            {
                "agent_slug": "my-agent",
                "tool_manifest_hash": "e3b0c442...",
                "tools": ["read_file", "send_email", "write_file"],
                "signature": "base64-signature...",
                "updated_at": "2024-01-15T10:30:00.000Z"
            }
        """
        slug = agent or self.agent
        if not slug:
            raise ValueError("Agent slug required")
        if not tools:
            raise ValueError("Tools list required")
        
        payload = {"tools": tools}
        if tool_descriptions:
            payload["tool_descriptions"] = tool_descriptions
        
        response = self._client.post(f"/v1/identity/{slug}/tools", json=payload)
        response.raise_for_status()
        return response.json()

    def get_tool_manifest(
        self,
        agent: Optional[str] = None,
    ) -> Dict[str, Any]:
        """
        Get the current tool manifest for an agent.
        
        Example:
            >>> manifest = ts.get_tool_manifest(agent="my-agent")
            >>> print(manifest["tools"])
            ["read_file", "write_file", "send_email"]
        
        Args:
            agent: Agent slug. Uses default agent if not provided.
            
        Returns:
            Tool manifest with hash, tools list, and signature.
        """
        slug = agent or self.agent
        if not slug:
            raise ValueError("Agent slug required")
        
        response = self._client.get(f"/v1/identity/{slug}/tools")
        response.raise_for_status()
        return response.json()

    def verify_tools(
        self,
        agent: Optional[str] = None,
        tools_used: Optional[list] = None,
    ) -> Dict[str, Any]:
        """
        Verify that tools used are in the agent's authorized manifest.
        
        Example:
            >>> result = ts.verify_tools(
            ...     agent="my-agent",
            ...     tools_used=["read_file", "delete_database"]
            ... )
            >>> print(result)
            {
                "verified": False,
                "authorized_tools": ["read_file"],
                "unauthorized_tools": ["delete_database"],
                "manifest_hash": "abc123..."
            }
        
        Args:
            agent: Agent slug. Uses default agent if not provided.
            tools_used: List of tools that were used.
            
        Returns:
            Verification result with authorized/unauthorized tools.
        """
        slug = agent or self.agent
        if not slug:
            raise ValueError("Agent slug required")
        if not tools_used:
            raise ValueError("tools_used list required")
        
        response = self._client.post(
            f"/v1/identity/{slug}/verify-tools",
            params={"tools_used": tools_used},
        )
        response.raise_for_status()
        return response.json()

    def register_voice_fingerprint(
        self,
        fingerprint_hash: str,
        agent: Optional[str] = None,
        provider: str = "custom",
        sample_count: int = 1,
    ) -> Dict[str, Any]:
        """
        Register a voice fingerprint for payment verification.
        
        Args:
            fingerprint_hash: Hash of voice embedding/fingerprint.
            agent: Agent slug. Uses default agent if not provided.
            provider: Voice ID provider (e.g., "elevenlabs", "resemble", "custom").
            sample_count: Number of voice samples used.
            
        Returns:
            Registration confirmation.
        """
        slug = agent or self.agent
        if not slug:
            raise ValueError("Agent slug required")
        
        payload = {
            "fingerprint_hash": fingerprint_hash,
            "provider": provider,
            "sample_count": sample_count,
        }
        
        response = self._client.post(f"/v1/identity/{slug}/voice-fingerprint", json=payload)
        response.raise_for_status()
        return response.json()

    def verify_voice(
        self,
        fingerprint_hash: str,
        agent: Optional[str] = None,
    ) -> Dict[str, Any]:
        """
        Verify a voice fingerprint for payment release.
        
        Args:
            fingerprint_hash: Hash of voice to verify.
            agent: Agent slug. Uses default agent if not provided.
            
        Returns:
            Verification result with confidence score and safe_for_payment flag.
        """
        slug = agent or self.agent
        if not slug:
            raise ValueError("Agent slug required")
        
        response = self._client.post(
            f"/v1/identity/{slug}/verify-voice",
            params={"fingerprint_hash": fingerprint_hash},
        )
        response.raise_for_status()
        return response.json()

    def close(self) -> None:
        """Close the HTTP client."""
        self._client.close()

    def __enter__(self) -> Treeship:
        return self

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