"""
NeuralDLP Client

Main client class for interacting with the NeuralDLP API.
"""
import httpx
from typing import Optional, Dict, Any, List
from datetime import datetime

from .models import (
    InspectionResult,
    TokenizationResult,
    RehydrationResult,
    Entity,
    MappingStatus
)
from .exceptions import (
    NeuralDLPError,
    AuthenticationError,
    RateLimitError,
    ValidationError,
    NotFoundError
)


class NeuralDLP:
    """
    NeuralDLP API Client
    
    Args:
        api_key: Your NeuralDLP API key
        base_url: API base URL (default: https://api.neuraldlp.com)
        timeout: Request timeout in seconds (default: 30)
        
    Example:
        >>> from neuraldlp import NeuralDLP
        >>> client = NeuralDLP(api_key="your_api_key")
        >>> result = client.inspect_input("My SSN is 123-45-6789")
        >>> print(result.risk_score)
        98
    """
    
    def __init__(
        self,
        api_key: str,
        base_url: str = "http://localhost:8000",  # TODO: Change to https://api.neuraldlp.com
        timeout: int = 30
    ):
        self.api_key = api_key
        self.base_url = base_url.rstrip("/")
        self.timeout = timeout
        
        self.client = httpx.Client(
            base_url=self.base_url,
            headers={
                "Authorization": f"Bearer {api_key}",
                "Content-Type": "application/json",
                "User-Agent": "neuraldlp-python/1.0.0"
            },
            timeout=timeout
        )
    
    def __enter__(self):
        return self
    
    def __exit__(self, *args):
        self.close()
    
    def close(self):
        """Close the HTTP client"""
        self.client.close()
    
    def _handle_error(self, response: httpx.Response):
        """Handle API error responses"""
        if response.status_code == 401:
            raise AuthenticationError("Invalid API key")
        elif response.status_code == 429:
            raise RateLimitError("Rate limit exceeded")
        elif response.status_code == 422:
            raise ValidationError(f"Validation error: {response.text}")
        elif response.status_code == 404:
            raise NotFoundError(response.text)
        elif response.status_code >= 500:
            raise NeuralDLPError(f"Server error: {response.status_code}")
        elif response.status_code >= 400:
            raise NeuralDLPError(f"Client error: {response.text}")
    
    def inspect_input(
        self,
        text: str,
        tenant_id: str = "default",
        user_id: Optional[str] = None,
        application: Optional[str] = None,
        metadata: Optional[Dict[str, Any]] = None
    ) -> InspectionResult:
        """
        Inspect input text for sensitive data and threats.
        
        Args:
            text: Text to inspect
            tenant_id: Tenant identifier (default: "default")
            user_id: User identifier (optional)
            application: Application name (optional)
            metadata: Additional metadata (optional)
            
        Returns:
            InspectionResult with detected entities, threats, and risk score
            
        Example:
            >>> result = client.inspect_input(
            ...     text="My email is john@example.com",
            ...     tenant_id="acme_corp",
            ...     user_id="user_123"
            ... )
            >>> print(f"Risk: {result.risk_score}, Action: {result.recommended_action}")
            Risk: 39, Action: ALLOW
        """
        payload = {
            "text": text,
            "context": {
                "tenant_id": tenant_id,
                "user_id": user_id,
                "application": application,
                "metadata": metadata
            }
        }
        
        response = self.client.post("/v1/inspect/input", json=payload)
        
        if response.status_code != 200:
            self._handle_error(response)
        
        return InspectionResult.from_dict(response.json())
    
    def inspect_output(
        self,
        text: str,
        mapping_id: Optional[str] = None,
        tenant_id: str = "default",
        user_id: Optional[str] = None
    ) -> InspectionResult:
        """
        Inspect AI output for data leakage.
        
        Args:
            text: Output text to inspect
            mapping_id: Original mapping ID from tokenization (optional)
            tenant_id: Tenant identifier
            user_id: User identifier (optional)
            
        Returns:
            InspectionResult with leakage detection results
            
        Example:
            >>> result = client.inspect_output(
            ...     text="The user PERSON_1 requested help",
            ...     mapping_id="map_abc123"
            ... )
            >>> print(result.leakage_detected)
            False
        """
        payload = {
            "text": text,
            "original_mapping_id": mapping_id,
            "context": {
                "tenant_id": tenant_id,
                "user_id": user_id
            }
        }
        
        response = self.client.post("/v1/inspect/output", json=payload)
        
        if response.status_code != 200:
            self._handle_error(response)
        
        return InspectionResult.from_dict(response.json())
    
    def tokenize(
        self,
        text: str,
        entities: List[Entity],
        tenant_id: str = "default",
        ttl_seconds: int = 3600,
        preserve_context: bool = True,
        deterministic: bool = True
    ) -> TokenizationResult:
        """
        Tokenize sensitive entities in text.
        
        Args:
            text: Text to tokenize
            entities: List of entities to tokenize
            tenant_id: Tenant identifier
            ttl_seconds: Time to live for mappings (default: 3600)
            preserve_context: Preserve surrounding context (default: True)
            deterministic: Use deterministic tokens (default: True)
            
        Returns:
            TokenizationResult with transformed text and mapping ID
            
        Example:
            >>> entities = [Entity(type="EMAIL", value="john@example.com", start=0, end=16)]
            >>> result = client.tokenize(
            ...     text="john@example.com needs help",
            ...     entities=entities
            ... )
            >>> print(result.transformed_text)
            EMAIL_abc123 needs help
        """
        payload = {
            "text": text,
            "entities": [e.to_dict() for e in entities],
            "context": {
                "tenant_id": tenant_id,
                "ttl_seconds": ttl_seconds
            },
            "options": {
                "preserve_context": preserve_context,
                "deterministic": deterministic
            }
        }
        
        response = self.client.post("/v1/transform/tokenize", json=payload)
        
        if response.status_code != 200:
            self._handle_error(response)
        
        return TokenizationResult.from_dict(response.json())
    
    def rehydrate(
        self,
        text: str,
        mapping_id: str,
        tenant_id: str = "default"
    ) -> RehydrationResult:
        """
        Rehydrate tokens back to original values.
        
        Args:
            text: Text with tokens
            mapping_id: Mapping ID from tokenization
            tenant_id: Tenant identifier
            
        Returns:
            RehydrationResult with restored text
            
        Example:
            >>> result = client.rehydrate(
            ...     text="EMAIL_abc123 needs help",
            ...     mapping_id="map_abc123"
            ... )
            >>> print(result.rehydrated_text)
            john@example.com needs help
        """
        payload = {
            "text": text,
            "mapping_id": mapping_id,
            "context": {
                "tenant_id": tenant_id
            }
        }
        
        response = self.client.post("/v1/transform/rehydrate", json=payload)
        
        if response.status_code != 200:
            self._handle_error(response)
        
        return RehydrationResult.from_dict(response.json())
    
    def inspect_and_tokenize(
        self,
        text: str,
        tenant_id: str = "default",
        user_id: Optional[str] = None,
        application: Optional[str] = None,
        auto_block: bool = True
    ) -> tuple[InspectionResult, Optional[TokenizationResult]]:
        """
        Convenience method: Inspect and automatically tokenize if needed.
        
        Args:
            text: Text to process
            tenant_id: Tenant identifier
            user_id: User identifier (optional)
            application: Application name (optional)
            auto_block: Raise error if action is BLOCK (default: True)
            
        Returns:
            Tuple of (InspectionResult, TokenizationResult or None)
            
        Raises:
            NeuralDLPError: If auto_block=True and action is BLOCK
            
        Example:
            >>> inspection, tokenization = client.inspect_and_tokenize(
            ...     text="My email is john@example.com"
            ... )
            >>> if tokenization:
            ...     print(tokenization.transformed_text)
        """
        # Step 1: Inspect
        inspection = self.inspect_input(
            text=text,
            tenant_id=tenant_id,
            user_id=user_id,
            application=application
        )
        
        # Step 2: Handle based on action
        if inspection.recommended_action == "BLOCK":
            if auto_block:
                raise NeuralDLPError(
                    f"Content blocked by policy. Risk score: {inspection.risk_score}"
                )
            return inspection, None
        
        elif inspection.recommended_action == "TOKENIZE":
            tokenization = self.tokenize(
                text=text,
                entities=inspection.entities,
                tenant_id=tenant_id
            )
            return inspection, tokenization
        
        else:  # ALLOW
            return inspection, None
    
    def get_mapping_status(
        self,
        mapping_id: str,
        tenant_id: str = "default"
    ) -> MappingStatus:
        """
        Get status of a mapping (TTL, expiration).
        
        Args:
            mapping_id: Mapping identifier
            tenant_id: Tenant identifier
            
        Returns:
            MappingStatus with TTL information
            
        Example:
            >>> status = client.get_mapping_status("map_abc123")
            >>> print(f"Expires in {status.ttl_seconds} seconds")
        """
        response = self.client.get(
            f"/v1/transform/mappings/{mapping_id}",
            params={"tenant_id": tenant_id}
        )
        
        if response.status_code != 200:
            self._handle_error(response)
        
        return MappingStatus.from_dict(response.json())
    
    def delete_mapping(
        self,
        mapping_id: str,
        tenant_id: str = "default"
    ) -> bool:
        """
        Delete a mapping before it expires.
        
        Args:
            mapping_id: Mapping identifier
            tenant_id: Tenant identifier
            
        Returns:
            True if deleted successfully
            
        Example:
            >>> client.delete_mapping("map_abc123")
            True
        """
        response = self.client.delete(
            f"/v1/transform/mappings/{mapping_id}",
            params={"tenant_id": tenant_id}
        )
        
        if response.status_code != 200:
            self._handle_error(response)
        
        return response.json().get("deleted", False)
    
    def health_check(self) -> Dict[str, Any]:
        """
        Check API health status.
        
        Returns:
            Health status dictionary
            
        Example:
            >>> status = client.health_check()
            >>> print(status['status'])
            healthy
        """
        response = self.client.get("/health")
        
        if response.status_code != 200:
            self._handle_error(response)
        
        return response.json()
