"""
Authevo ZK Agent

Zero-Knowledge agent that keeps sensitive data local and only sends
cryptographic hashes to Authevo Cloud.

Authevo never sees your payload data - only the hash.
"""

import datetime
import hashlib
import uuid
import httpx
from typing import Dict, Any, Optional
from dataclasses import dataclass

from .crypto import KeyPair
from .client import AuthevoClient
from .canonicalize import canonicalize
from .config import AuthevoConfig
from .local_store import LocalActionStore


@dataclass
class ZKActionResult:
    """Result from a ZK action execution."""
    action_id: str               # Cloud action ID
    local_id: str                # Local storage ID
    action_hash: str             # SHA256 hash
    verified: bool               # Cloud verified signature
    transparency_log_idx: Optional[int] = None  # Transparency log index


class AuthevoZKAgent:
    """
    Zero-Knowledge Authevo Agent.
    
    This agent:
    1. Computes action hash locally
    2. Signs the hash with Ed25519
    3. Stores full action in LOCAL storage (your infrastructure)
    4. Sends ONLY hash + signature to Authevo Cloud
    
    Authevo Cloud never sees your payload data.
    
    Example:
        agent = AuthevoZKAgent.generate()
        
        result = await agent.execute_action(
            action="refund",
            payload={"amount": 50, "customer": "john@example.com"}
        )
        
        # Your data stays local
        # Authevo only has: hash + signature
    """
    
    def __init__(
        self, 
        keypair: KeyPair,
        config: Optional[AuthevoConfig] = None,
        local_store: Optional[LocalActionStore] = None,
        policy_path: str = "authevo/refund/allow"
    ):
        """
        Initialize ZK Agent.
        
        Args:
            keypair: Ed25519 keypair for signing
            config: Authevo configuration
            local_store: Custom local storage (defaults to SQLite)
            policy_path: OPA policy path for evaluation
        """
        self.keypair = keypair
        self.config = config or AuthevoConfig.from_env()
        self.client = AuthevoClient(config=self.config)
        self.local_store = local_store or LocalActionStore()
        self.policy_path = policy_path
        self._http_client = httpx.AsyncClient(timeout=self.config.timeout)
    
    @property
    def did(self) -> str:
        """Agent's DID (did:key:...)"""
        return self.keypair.did
    
    @classmethod
    def generate(cls, config: Optional[AuthevoConfig] = None, **kwargs):
        """Generate a new agent with a random keypair."""
        return cls(KeyPair.generate(), config=config, **kwargs)
    
    @classmethod
    def load(cls, private_key_b64: str, config: Optional[AuthevoConfig] = None, **kwargs):
        """Load an agent from a private key."""
        return cls(KeyPair.from_private_key(private_key_b64), config=config, **kwargs)
    
    async def close(self):
        """Close all HTTP connections."""
        await self.client.close()
        await self._http_client.aclose()
    
    async def __aenter__(self):
        return self
    
    async def __aexit__(self, exc_type, exc_val, exc_tb):
        await self.close()
        return False
    
    async def evaluate_policy(self, action: str, payload: Dict[str, Any]) -> Dict[str, Any]:
        """Evaluate OPA policy for the action."""
        if not self.config.opa_url:
            # No OPA configured, default to permit
            return {
                "decision": "Permit",
                "reason": "No policy configured",
                "policyId": self.policy_path
            }
        
        opa_input = {
            "input": {
                "action": action,
                "payload": payload,
                "agent": self.keypair.did
            }
        }
        
        try:
            response = await self._http_client.post(
                f"{self.config.opa_url}/v1/data/{self.policy_path}",
                json=opa_input
            )
            response.raise_for_status()
            result = response.json().get("result", {})
            
            if isinstance(result, bool):
                decision = "Permit" if result else "Deny"
                reason = "Policy allowed" if result else "Policy denied"
            else:
                allowed = result.get("allow", False) if isinstance(result, dict) else result
                decision = "Permit" if allowed else "Deny"
                reason = result.get("reason", "Policy evaluation") if isinstance(result, dict) else "Policy evaluation"
            
            return {
                "decision": decision,
                "reason": reason,
                "policyId": self.policy_path
            }
        except Exception as e:
            return {
                "decision": "Deny",
                "reason": f"Policy evaluation failed: {str(e)}",
                "policyId": self.policy_path
            }
    
    async def execute_action(
        self, 
        action: str, 
        payload: Dict[str, Any],
        evaluate_policy: bool = True
    ) -> ZKActionResult:
        """
        Execute a Zero-Knowledge verifiable action.
        
        1. Optionally evaluates OPA policy
        2. Builds full action object
        3. Computes SHA256 hash
        4. Signs the hash
        5. Stores full action LOCALLY
        6. Sends only hash + signature to Authevo Cloud
        
        Args:
            action: Action type (e.g., "refund", "approve")
            payload: Action payload (stays local!)
            evaluate_policy: Whether to check OPA policy first
            
        Returns:
            ZKActionResult with action IDs and verification status
        """
        # 1. Evaluate policy (optional)
        policy_decision = None
        if evaluate_policy:
            policy_decision = await self.evaluate_policy(action, payload)
        
        # 2. Build full action object
        timestamp = datetime.datetime.now(datetime.timezone.utc).isoformat().replace("+00:00", "Z")
        nonce = str(uuid.uuid4())
        
        full_action = {
            "agent": self.keypair.did,
            "action": action,
            "payload": payload,
            "timestamp": timestamp,
            "nonce": nonce
        }
        
        if policy_decision:
            full_action["policyDecision"] = policy_decision
        
        # 3. Canonicalize and compute hash
        canonical_bytes = canonicalize(full_action)
        action_hash = hashlib.sha256(canonical_bytes).hexdigest()
        
        # 4. Sign the hash (not the full action)
        hash_bytes = bytes.fromhex(action_hash)
        signature = self.keypair.sign(hash_bytes)
        
        # 5. Store full action LOCALLY
        local_id = self.local_store.save(full_action, action_hash)
        
        # 6. Send only hash + signature to cloud
        try:
            response = await self._submit_zk_action(
                action_hash=action_hash,
                signature=signature,
                timestamp=timestamp,
                action_type=action,
                nonce=nonce
            )
            
            # Update local record with cloud response
            self.local_store.update_cloud_response(
                action_hash=action_hash,
                cloud_action_id=response.get("actionId"),
                transparency_log_idx=response.get("transparencyLogIndex")
            )
            
            return ZKActionResult(
                action_id=response.get("actionId", ""),
                local_id=local_id,
                action_hash=action_hash,
                verified=response.get("verified", False),
                transparency_log_idx=response.get("transparencyLogIndex")
            )
            
        except Exception as e:
            # Action is stored locally even if cloud submission fails
            # Customer can retry or operate offline
            raise RuntimeError(
                f"Cloud submission failed (action saved locally as {local_id}): {str(e)}"
            )
    
    async def _submit_zk_action(
        self,
        action_hash: str,
        signature: str,
        timestamp: str,
        action_type: str,
        nonce: str
    ) -> Dict[str, Any]:
        """Submit ZK action to Authevo Cloud."""
        response = await self._http_client.post(
            f"{self.config.api_url}/actions",
            json={
                "actionHash": action_hash,
                "signature": signature,
                "agentDID": self.keypair.did,
                "timestamp": timestamp,
                "actionType": action_type,
                "nonce": nonce
            },
            headers={
                "Authorization": f"Bearer {self.config.api_key}" if self.config.api_key else "",
                "Content-Type": "application/json"
            }
        )
        
        if response.status_code >= 400:
            error_data = response.json() if response.content else {}
            raise RuntimeError(
                f"API error ({response.status_code}): {error_data.get('error', 'Unknown error')}"
            )
        
        return response.json()
    
    async def verify_action(self, action_hash: str) -> Dict[str, Any]:
        """
        Verify an action exists in Authevo Cloud.
        
        Returns cloud attestation for the hash.
        """
        response = await self._http_client.get(
            f"{self.config.api_url}/actions/verify/{action_hash}",
            headers={
                "Authorization": f"Bearer {self.config.api_key}" if self.config.api_key else ""
            }
        )
        
        return response.json()
    
    def get_local_action(self, action_hash: str) -> Optional[Dict[str, Any]]:
        """
        Get full action data from local storage.
        
        Use this for dispute resolution - provide raw data that can be
        verified against the hash stored in Authevo Cloud.
        """
        return self.local_store.get_by_hash(action_hash)
    
    def export_for_audit(
        self, 
        start_date: datetime.datetime, 
        end_date: datetime.datetime
    ):
        """Export local actions for external audit."""
        return self.local_store.export_for_audit(start_date, end_date, self.keypair.did)
