import requests
import time
from typing import List, Optional, Dict, Any, Union
from .models import ContextParams, AddressOutput, NameOutput, MoneyOutput
from .exceptions import (
    ETLDAuthenticationError,
    ETLDRateLimitError,
    ETLDProcessingError,
    ETLDValidationError
)

class ETLDClient:
    def __init__(self, api_key: str, base_url: str = "https://api.etl-d.net"):
        self.api_key = api_key
        self.base_url = base_url.rstrip("/")
        self.session = requests.Session()
        self.session.headers.update({
            "X-API-KEY": self.api_key,
            "Content-Type": "application/json",
            "User-Agent": "etld-python-sdk/0.1.0"
        })

    def _handle_response(self, response: requests.Response) -> Dict[str, Any]:
        if response.status_code == 401 or response.status_code == 403:
            raise ETLDAuthenticationError(f"Invalid API Key: {response.text}")
        elif response.status_code == 429:
            raise ETLDRateLimitError(f"Rate limit exceeded: {response.text}")
        elif response.status_code == 422:
            raise ETLDValidationError(f"Validation Error: {response.json().get('detail')}")
        
        try:
            response.raise_for_status()
        except requests.exceptions.HTTPError as e:
            raise ETLDProcessingError(f"API Request failed: {str(e)}")
            
        return response.json()

    def enrich_address(self, address: str, context: Optional[Union[Dict[str, Any], ContextParams]] = None) -> AddressOutput:
        if isinstance(context, ContextParams):
            context = context.to_dict()
        
        payload = {
            "address_line": address,
            "context": context or {}
        }
        
        data = self._handle_response(self.session.post(f"{self.base_url}/v1/enrich/address", json=payload))
        return AddressOutput(**data)

    def enrich_name(self, name: str, context: Optional[Union[Dict[str, Any], ContextParams]] = None) -> NameOutput:
        if isinstance(context, ContextParams):
            context = context.to_dict()
            
        payload = {
            "full_name": name,
            "context": context or {}
        }
        
        data = self._handle_response(self.session.post(f"{self.base_url}/v1/enrich/name", json=payload))
        return NameOutput(**data)

    def enrich_amount(self, amount: str, context: Optional[Union[Dict[str, Any], ContextParams]] = None) -> MoneyOutput:
        if isinstance(context, ContextParams):
            context = context.to_dict()
            
        payload = {
            "amount_string": amount,
            "context": context or {}
        }
        
        data = self._handle_response(self.session.post(f"{self.base_url}/v1/enrich/amount", json=payload))
        return MoneyOutput(**data)

    def batch_process(
        self, 
        items: List[str], 
        entity_type: str, 
        context: Optional[Union[Dict[str, Any], ContextParams]] = None,
        poll_interval: int = 2,
        max_retries: int = 30
    ) -> List[Any]:
        if isinstance(context, ContextParams):
            context = context.to_dict()
            
        payload = {
            "items": items,
            "context": context or {}
        }
        
        # 1. Start Batch
        start_data = self._handle_response(
            self.session.post(f"{self.base_url}/v1/enrich/batch/{entity_type}", json=payload)
        )
        task_id = start_data.get("task_id")
        
        # 2. Poll
        retries = 0
        while retries < max_retries:
            status_data = self._handle_response(
                self.session.get(f"{self.base_url}/v1/enrich/batch/{task_id}")
            )
            
            status = status_data.get("status")
            if status == "SUCCESS":
                return status_data.get("results", [])
            elif status == "FAILURE":
                raise ETLDProcessingError(f"Batch processing failed: {status_data.get('error')}")
            
            time.sleep(poll_interval)
            retries += 1
            
        raise ETLDProcessingError(f"Batch processing timed out for task {task_id}")
