from __future__ import annotations
import httpx
from .models import Company, ValidationResult, CountriesResponse, HealthResponse, BatchResponse
from .errors import VeriCorpError, RateLimitError, InvalidTaxIdError, NotFoundError, TimeoutError


class VeriCorp:
    """Synchronous client for the VeriCorp API."""

    BASE_URL = "https://vericorp-api.p.rapidapi.com"

    def __init__(self, api_key: str, *, timeout: float = 10.0, max_retries: int = 2):
        self._api_key = api_key
        self._timeout = timeout
        self._max_retries = max_retries
        self._client = httpx.Client(
            base_url=self.BASE_URL,
            headers={
                "X-RapidAPI-Key": api_key,
                "X-RapidAPI-Host": "vericorp-api.p.rapidapi.com",
            },
            timeout=timeout,
        )

    def _request(self, method: str, path: str, **kwargs) -> httpx.Response:
        last_error: Exception | None = None
        for attempt in range(self._max_retries + 1):
            try:
                response = self._client.request(method, path, **kwargs)

                if response.status_code == 429:
                    retry_after = response.headers.get("Retry-After")
                    if attempt < self._max_retries:
                        import time
                        wait = int(retry_after) if retry_after else (2 ** attempt)
                        time.sleep(wait)
                        continue
                    raise RateLimitError(int(retry_after) if retry_after else None)

                if response.status_code == 503 and attempt < self._max_retries:
                    import time
                    time.sleep(2 ** attempt)
                    continue

                return response
            except httpx.TimeoutException:
                last_error = TimeoutError()
                if attempt < self._max_retries:
                    continue
                raise last_error
            except (RateLimitError, TimeoutError):
                raise
            except httpx.HTTPError as e:
                last_error = VeriCorpError(str(e))
                if attempt < self._max_retries:
                    continue
                raise last_error

        raise last_error or VeriCorpError("Request failed")

    def _handle_error(self, response: httpx.Response, tax_id: str | None = None) -> None:
        if response.status_code == 400:
            raise InvalidTaxIdError(tax_id or "unknown")
        if response.status_code == 404:
            raise NotFoundError(tax_id or "unknown")
        if response.status_code == 429:
            retry_after = response.headers.get("Retry-After")
            raise RateLimitError(int(retry_after) if retry_after else None)
        if response.status_code >= 400:
            raise VeriCorpError(f"API error: {response.status_code}", response.status_code)

    def lookup(self, tax_id: str, *, force_refresh: bool = False) -> Company:
        params = {"force_refresh": "true"} if force_refresh else {}
        response = self._request("GET", f"/v1/company/{tax_id}", params=params)
        if not response.is_success:
            self._handle_error(response, tax_id)
        return Company.model_validate(response.json())

    def lookup_gb(self, company_number: str, *, force_refresh: bool = False) -> Company:
        params = {"force_refresh": "true"} if force_refresh else {}
        response = self._request("GET", f"/v1/company/gb/{company_number}", params=params)
        if not response.is_success:
            self._handle_error(response, company_number)
        return Company.model_validate(response.json())

    def validate(self, tax_id: str) -> ValidationResult:
        response = self._request("GET", f"/v1/validate/{tax_id}")
        if not response.is_success:
            self._handle_error(response, tax_id)
        return ValidationResult.model_validate(response.json())

    def batch(self, tax_ids: list[str]) -> BatchResponse:
        response = self._request("POST", "/v1/batch", json={"tax_ids": tax_ids})
        if not response.is_success:
            self._handle_error(response)
        return BatchResponse.model_validate(response.json())

    def countries(self) -> CountriesResponse:
        response = self._request("GET", "/v1/countries")
        if not response.is_success:
            self._handle_error(response)
        return CountriesResponse.model_validate(response.json())

    def health(self) -> HealthResponse:
        response = self._request("GET", "/v1/health")
        if not response.is_success:
            self._handle_error(response)
        return HealthResponse.model_validate(response.json())

    def close(self) -> None:
        self._client.close()

    def __enter__(self):
        return self

    def __exit__(self, *args):
        self.close()
