"""
Clox Client SDK

A simple Python client for interacting with the Clox Document Processing API.

Example usage:
    from clox_client import Clox

    client = Clox(
        api_url="http://localhost:8000",
        api_key="your_api_key"
    )

    # Process a document
    result = client.process_document("invoice.pdf")
    print(result)
"""

import os
from typing import Optional, Dict, Any, List, BinaryIO, Union
from pathlib import Path
import requests
from requests.adapters import HTTPAdapter
from urllib3.util.retry import Retry
import logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)    

class CloxException(Exception):
    """Base exception for Clox client errors"""
    pass


class CloxAuthenticationError(CloxException):
    """Raised when authentication fails"""
    pass


class CloxAPIError(CloxException):
    """Raised when API request fails"""
    pass


class Clox:
    """
    Clox Client SDK for document processing and fraud detection.

    This client provides a simple interface to interact with the Clox API.

    Args:
        api_url: Base URL of the Clox API (default: http://localhost:8000)
        api_key: API key for authentication
        timeout: Request timeout in seconds (default: 30)
        max_retries: Maximum number of retries for failed requests (default: 3)

    Example:
        >>> client = Clox(
        ...     api_url="http://localhost:8000",
        ...     api_key="your_api_key"
        ... )
        >>> result = client.process_document("invoice.pdf")
        >>> print(result["document_id"])
    """

    def __init__(
        self,
        api_url: Optional[str] = None,
        api_key: Optional[str] = None,
        timeout: int = 30,
        max_retries: int = 3
    ):
        self.api_url = (api_url or os.getenv("CLOX_API_URL", "http://localhost:8000")).rstrip("/")
        self.api_key = api_key or os.getenv("CLOX_API_KEY")
        self.timeout = timeout

        if not self.api_key:
            raise CloxException(
                "api_key is required. "
                "Provide it as an argument or set CLOX_API_KEY environment variable."
            )

        # Setup session with retry logic
        self.session = requests.Session()
        retry_strategy = Retry(
            total=max_retries,
            backoff_factor=1,
            status_forcelist=[429, 500, 502, 503, 504],
            allowed_methods=["HEAD", "GET", "PUT", "DELETE", "OPTIONS", "TRACE", "POST"]
        )
        adapter = HTTPAdapter(max_retries=retry_strategy)
        self.session.mount("http://", adapter)
        self.session.mount("https://", adapter)

    def _get_headers(self) -> Dict[str, str]:
        """Get headers with API key authentication."""
        return {
            "X-API-KEY": self.api_key
        }

    def _make_request(
        self,
        method: str,
        endpoint: str,
        **kwargs
    ) -> Union[Dict[str, Any], bytes]:
        """Make an authenticated request to the API."""
        url = f"{self.api_url}{endpoint}"
        headers = self._get_headers()

        if "headers" in kwargs:
            headers.update(kwargs.pop("headers"))

        try:
            response = self.session.request(
                method=method,
                url=url,
                headers=headers,
                timeout=self.timeout,
                **kwargs
            )

            # Handle 401 (unauthorized) - API key may be invalid
            if response.status_code == 401:
                raise CloxAuthenticationError(
                    "Authentication failed. Please check your API key."
                )

            response.raise_for_status()

            # Parse response based on Content-Type header
            content_type = response.headers.get('Content-Type', '').lower()

            if 'application/json' in content_type:
                return response.json()
            elif 'application/octet-stream' in content_type or \
                 'application/pdf' in content_type or \
                 'image/' in content_type:
                # Return binary content for documents/images
                return response.content
            else:
                # Default to JSON for backward compatibility
                try:
                    return response.json()
                except (ValueError, requests.exceptions.JSONDecodeError):
                    # If JSON parsing fails, return binary content
                    return response.content

        except requests.exceptions.HTTPError as e:
            error_msg = f"API request failed: {str(e)}"
            try:
                error_data = e.response.json()
                error_msg = f"{error_msg} - {error_data.get('detail', '')}"
            except:
                pass
            raise CloxAPIError(error_msg)
        except requests.exceptions.RequestException as e:
            raise CloxAPIError(f"Request failed: {str(e)}")

    # ==================== Document Processing ====================

    def process_document(
        self,
        file: Union[str, Path, BinaryIO],
        filename: Optional[str] = None,
        **kwargs
    ) -> Dict[str, Any]:
        """
        Process a document (PDF or image) with OCR and fraud detection.

        Args:
            file: Path to the file or file-like object
            filename: Optional filename (used when file is file-like object)
            **kwargs: Additional parameters to pass to the API

        Returns:
            Dictionary containing processing results with keys:
                - document_id: Unique document identifier
                - status: Processing status
                - text: Extracted text (if OCR enabled)
                - fraud_detection: Fraud detection results (if enabled)
                - metadata: Document metadata

        Example:
            >>> result = client.process_document("invoice.pdf")
            >>> print(f"Document ID: {result['document_id']}")
            >>> print(f"Extracted text: {result['text']}")
        """
        if isinstance(file, (str, Path)):
            file = Path(file)
            if not file.exists():
                raise CloxException(f"File not found: {file}")

            with open(file, "rb") as f: 
                files = [('file', (file.name, f, self._get_content_type(file)))]
                return self._make_request("POST", "/v1/process_document", files=files, data=kwargs)
        else:
            # file is a file-like object
            if not filename:
                filename = "document"
            files = [('file', (filename, file, self._get_content_type(file)))]
            return self._make_request("POST", "/v1/process_document", files=files, data=kwargs)

    def process_documents(
        self,
        files: List[Union[str, Path, BinaryIO]],
        **kwargs
    ) -> Dict[str, Any]:
        """
        Process multiple documents (PDF or image) with OCR and fraud detection.

        Args:
            files: List of file paths (strings or Path objects), file-like objects, or binary data
            **kwargs: Additional parameters to pass to the API

        Returns:
            Dictionary containing processing results for all documents

        Example:
            >>> result = client.process_documents([
            ...     "document1.pdf",
            ...     "document2.jpg",
            ...     open("document3.pdf", "rb")
            ... ])
        """
        # Open all files and create the files list in the correct format
        file_objects = []
        opened_files = []  # Track files we open so we can close them later

        try:
            for file in files:
                if isinstance(file, (str, Path)):
                    file = Path(file)
                    if not file.exists():
                        raise CloxException(f"File not found: {file}")
                    f = open(file, "rb")
                    opened_files.append(f)
                    file_objects.append(('files', (file.name, f, self._get_content_type(file))))
                elif isinstance(file, BinaryIO):
                    file_objects.append(('files', (file.name, file, self._get_content_type(file))))
                else:
                    raise CloxException(f"Invalid file type: {type(file)}")

            return self._make_request("POST", "/v1/process_documents", files=file_objects, data=kwargs)
        finally:
            # Close any files we opened
            for f in opened_files:
                f.close()

    def get_document(self, document_id: str) -> bytes:
        """
        Get document by ID (returns binary content).

        Args:
            document_id: Document identifier

        Returns:
            Binary content of the document (PDF, image, etc.)

        Example:
            >>> doc_bytes = client.get_document("doc_123")
            >>> with open("output.pdf", "wb") as f:
            ...     f.write(doc_bytes)
        """
        return self._make_request("GET", f"/v1/document/{document_id}")

    def get_annotated_document(self, document_id: str) -> bytes:
        """
        Get annotated document by ID (returns binary content).

        Args:
            document_id: Document identifier

        Returns:
            Binary content of the annotated document (PDF, image, etc.)

        Example:
            >>> annotated_bytes = client.get_annotated_document("doc_123")
            >>> with open("annotated.pdf", "wb") as f:
            ...     f.write(annotated_bytes)
        """
        return self._make_request("GET", f"/v1/annotated_document/{document_id}")

    def list_documents(
        self,
        page: int = 1,
        limit: int = 10,
        batch_id: Optional[str] = None
    ) -> Dict[str, Any]:
        """
        List documents with pagination.

        Args:
            page: Page number (default: 1)
            limit: Results per page (default: 10)
            batch_id: Filter by batch ID

        Returns:
            Dictionary containing paginated document list

        Example:
            >>> docs = client.list_documents(page=1, limit=20, batch_id="batch_123")
            >>> for doc in docs["documents"]:
            ...     print(doc["document_id"])
        """
        params = {"page": page, "limit": limit}
        if batch_id:
            params["batch_id"] = batch_id
        return self._make_request("GET", "/v1/client_documents", params=params)

    # ==================== Batch Processing ====================

    def get_batch_details(self, batch_id: str) -> Dict[str, Any]:
        """
        Get detailed batch information.

        Args:
            batch_id: Batch identifier

        Returns:
            Dictionary with batch details
        """
        return self._make_request("GET", "/v1/batch_details", params={"batch_id": batch_id})

    def get_batch_requests(self, page: int = 1, limit: int = 10) -> List[Dict[str, Any]]:
        """
        Get all batch requests.

        Args:
            page: Page number (default: 1)
            limit: Results per page (default: 10)

        Returns:
            List of batch request dictionaries
        """
        params = {"page": page, "page_size": limit}
        return self._make_request("GET", "/v1/batch_requests", params=params)

    # ==================== User Management ====================

    def create_user(self, user_data: Dict[str, Any]) -> Dict[str, Any]:
        """
        Create a new user.

        Args:
            user_data: Dictionary containing user information

        Returns:
            Dictionary containing created user details

        Example:
            >>> user = client.create_user({
            ...     "email": "user@example.com",
            ...     "first_name": "John",
            ...     "last_name": "Doe",
            ...     "password": "password",
            ...     "invitation_code": "INV123",
            ...     "phone": "+1234567890",
            ... })
        """
        return self._make_request("POST", "/v1/register", json=user_data)

    def list_users(self) -> List[Dict[str, Any]]:
        """
        List all users.

        Returns:
            List of user dictionaries
        """
        return self._make_request("GET", "/v1/users")

    def update_user(self, user_id: str, user_data: Dict[str, Any]) -> Dict[str, Any]:
        """
        Update user information.

        Args:
            user_id: User identifier
            user_data: Dictionary containing updated user information

        Returns:
            Dictionary containing updated user details
        """
        return self._make_request("PUT", f"/v1/user/{user_id}", json=user_data)

    def delete_user(self, user_id: str) -> Dict[str, Any]:
        """
        Delete a user.

        Args:
            user_id: User identifier
        """
        return self._make_request("DELETE", f"/v1/user/{user_id}")
    # ==================== Dashboard & Stats ====================

    def get_dashboard_stats(self) -> Dict[str, Any]:
        """
        Get dashboard statistics.

        Returns:
            Dictionary containing statistics (document counts, processing stats, etc.)

        Example:
            >>> stats = client.get_dashboard_stats()
            >>> print(f"Total documents: {stats['total_documents']}")
        """
        return self._make_request("GET", "/v1/dashboard/stats")

    # ==================== Health Check ====================

    def health_check(self) -> Dict[str, Any]:
        """
        Check API health status.

        Returns:
            Dictionary containing health status

        Example:
            >>> health = client.health_check()
            >>> print(health["status"])
        """
        response = self.session.get(f"{self.api_url}/health", timeout=self.timeout)
        response.raise_for_status()
        return response.json()

    # ==================== Helper Methods ====================

    @staticmethod
    def _get_content_type(file: Union[Path, BinaryIO]) -> str:
        """Determine content type based on file extension."""
        if isinstance(file, Path):
            suffix = file.suffix.lower()
        else:
            suffix = file.name.split(".")[-1].lower()
        content_types = {
            ".pdf": "application/pdf",
            ".png": "image/png",
            ".jpg": "image/jpeg",
            ".jpeg": "image/jpeg",
            ".tiff": "image/tiff",
            ".tif": "image/tiff",
        }
        return content_types.get(suffix, "application/octet-stream")
    def close(self):
        """Close the session."""
        self.session.close()

    def __enter__(self):
        """Context manager entry."""
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """Context manager exit."""
        self.close()

    def __repr__(self):
        return f"Clox(api_url='{self.api_url}', api_key='***')"


# Convenience function for quick usage
def create_client(
    api_url: Optional[str] = None,
    api_key: Optional[str] = None
) -> Clox:
    """
    Create a Clox client instance.

    Args:
        api_url: Base URL of the Clox API
        api_key: API key for authentication

    Returns:
        Clox client instance

    Example:
        >>> from clox_client import create_client
        >>> client = create_client(api_key="your_api_key")
        >>> result = client.process_document("invoice.pdf")
    """
    return Clox(api_url=api_url, api_key=api_key)
