"""
Ashr Labs API Client.

This module provides the main client for interacting with the Ashr Labs API.
"""

import json
from typing import Any
from urllib.request import Request, urlopen
from urllib.error import HTTPError, URLError

from .exceptions import (
    AshrLabsError,
    AuthenticationError,
    AuthorizationError,
    NotFoundError,
    ValidationError,
    RateLimitError,
    ServerError,
)
from .models import Dataset, Run, Request as RequestModel, APIKey, Session


class AshrLabsClient:
    """
    Client for interacting with the Ashr Labs API.

    This client supports all API key accessible endpoints:
    - Datasets: get, list
    - Runs: create, delete
    - Requests: create, get
    - API Keys: list, revoke

    Example:
        ```python
        from ashr_labs import AshrLabsClient

        client = AshrLabsClient(
            api_key="tp_your_api_key_here",
            base_url="https://api.ashr.io/testing-platform-api"
        )

        # List datasets
        datasets = client.list_datasets(tenant_id=1)

        # Create a run
        run = client.create_run(
            tenant_id=1,
            dataset_id=42,
            result={"score": 0.95, "status": "passed"}
        )
        ```
    """

    def __init__(
        self,
        api_key: str,
        base_url: str,
        timeout: int = 30,
    ):
        """
        Initialize the Ashr Labs client.

        Args:
            api_key: Your Ashr Labs API key (starts with 'tp_').
            base_url: The base URL of the API (e.g., 'https://api.ashr.io/testing-platform-api').
            timeout: Request timeout in seconds (default: 30).

        Raises:
            ValueError: If the API key format is invalid.
        """
        if not api_key or not api_key.startswith("tp_"):
            raise ValueError("Invalid API key format. API keys must start with 'tp_'")

        self.api_key = api_key
        self.base_url = base_url.rstrip("/")
        self.timeout = timeout

    def _make_request(self, function: str, **params: Any) -> dict:
        """
        Make an authenticated request to the API.

        Args:
            function: The API function to call.
            **params: Additional parameters for the function.

        Returns:
            The JSON response from the API.

        Raises:
            AshrLabsError: If the request fails.
        """
        payload = {"function": function, **params}

        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {self.api_key}",
        }

        data = json.dumps(payload).encode("utf-8")
        request = Request(self.base_url, data=data, headers=headers, method="POST")

        try:
            with urlopen(request, timeout=self.timeout) as response:
                response_data = json.loads(response.read().decode("utf-8"))
        except HTTPError as e:
            status_code = e.code
            try:
                error_body = json.loads(e.read().decode("utf-8"))
                message = error_body.get("message", str(e))
            except (json.JSONDecodeError, AttributeError):
                message = str(e)

            self._raise_for_status(status_code, message, error_body if "error_body" in dir() else None)
        except URLError as e:
            raise AshrLabsError(f"Network error: {e.reason}")
        except TimeoutError:
            raise AshrLabsError("Request timed out")

        # Check for API-level errors
        if response_data.get("status") == "error":
            message = response_data.get("message", "Unknown error")
            self._raise_for_status(400, message, response_data)

        return response_data

    def _raise_for_status(
        self, status_code: int, message: str, response: dict | None = None
    ) -> None:
        """Raise the appropriate exception based on status code."""
        if status_code == 401:
            raise AuthenticationError(message, status_code, response)
        elif status_code == 403:
            raise AuthorizationError(message, status_code, response)
        elif status_code == 404:
            raise NotFoundError(message, status_code, response)
        elif status_code == 422:
            raise ValidationError(message, status_code, response)
        elif status_code == 429:
            raise RateLimitError(message, status_code, response)
        elif status_code >= 500:
            raise ServerError(message, status_code, response)
        else:
            raise AshrLabsError(message, status_code, response)

    # =========================================================================
    # Dataset Operations
    # =========================================================================

    def get_dataset(
        self,
        dataset_id: int,
        include_signed_urls: bool = False,
        url_expires_seconds: int = 3600,
    ) -> Dataset:
        """
        Retrieve a dataset by ID.

        Args:
            dataset_id: The ID of the dataset to retrieve.
            include_signed_urls: Whether to include signed S3 URLs for media files.
            url_expires_seconds: Expiration time for signed URLs in seconds.

        Returns:
            The dataset object.

        Raises:
            NotFoundError: If the dataset is not found.
            AuthorizationError: If you don't have access to this dataset.
        """
        response = self._make_request(
            "get_dataset",
            dataset_id=dataset_id,
            include_signed_urls=include_signed_urls,
            url_expires_seconds=url_expires_seconds,
        )
        return response["dataset"]

    def list_datasets(
        self,
        tenant_id: int,
        limit: int = 50,
        offset: int = 0,
        include_signed_urls: bool = False,
        url_expires_seconds: int = 3600,
    ) -> dict:
        """
        List datasets for a tenant.

        Args:
            tenant_id: The ID of the tenant.
            limit: Maximum number of datasets to return (default: 50).
            offset: Number of datasets to skip for pagination.
            include_signed_urls: Whether to include signed S3 URLs for media files.
            url_expires_seconds: Expiration time for signed URLs in seconds.

        Returns:
            A dict containing 'datasets' list and pagination info.
        """
        return self._make_request(
            "list_datasets",
            tenant_id=tenant_id,
            limit=limit,
            offset=offset,
            include_signed_urls=include_signed_urls,
            url_expires_seconds=url_expires_seconds,
        )

    # =========================================================================
    # Run Operations
    # =========================================================================

    def create_run(
        self,
        tenant_id: int,
        dataset_id: int,
        result: dict[str, Any],
        runner_id: int | None = None,
    ) -> Run:
        """
        Create a new test run.

        Args:
            tenant_id: The ID of the tenant.
            dataset_id: The ID of the dataset this run is for.
            result: The run result data (metrics, status, etc.).
            runner_id: Optional ID of the user who ran the test.

        Returns:
            The created run object.

        Example:
            ```python
            run = client.create_run(
                tenant_id=1,
                dataset_id=42,
                result={
                    "score": 0.95,
                    "status": "passed",
                    "metrics": {"accuracy": 0.98, "latency_ms": 150}
                }
            )
            ```
        """
        params = {
            "tenant_id": tenant_id,
            "dataset_id": dataset_id,
            "result": result,
        }
        if runner_id is not None:
            params["runner_id"] = runner_id

        response = self._make_request("create_run", **params)
        return response["run"]

    def delete_run(self, run_id: int) -> dict:
        """
        Delete a test run.

        Args:
            run_id: The ID of the run to delete.

        Returns:
            Confirmation of deletion.

        Raises:
            NotFoundError: If the run is not found.
        """
        return self._make_request("delete_run", run_id=run_id)

    def get_run(self, run_id: int) -> Run:
        """
        Retrieve a run by ID.

        Args:
            run_id: The ID of the run to retrieve.

        Returns:
            The run object.

        Raises:
            NotFoundError: If the run is not found.
        """
        response = self._make_request("get_run", run_id=run_id)
        return response["run"]

    def list_runs(
        self,
        tenant_id: int | None = None,
        dataset_id: int | None = None,
        limit: int = 50,
        offset: int = 0,
    ) -> dict:
        """
        List runs for a tenant or dataset.

        Args:
            tenant_id: Filter by tenant ID.
            dataset_id: Filter by dataset ID.
            limit: Maximum number of runs to return.
            offset: Number of runs to skip for pagination.

        Returns:
            A dict containing 'runs' list and pagination info.
        """
        params = {"limit": limit, "offset": offset}
        if tenant_id is not None:
            params["tenant_id"] = tenant_id
        if dataset_id is not None:
            params["dataset_id"] = dataset_id

        return self._make_request("list_runs", **params)

    # =========================================================================
    # Request Operations
    # =========================================================================

    def create_request(
        self,
        tenant_id: int,
        requestor_id: int,
        request_name: str,
        request: dict[str, Any],
        request_input_schema: dict[str, Any] | None = None,
    ) -> RequestModel:
        """
        Create a new request.

        Args:
            tenant_id: The ID of the tenant.
            requestor_id: The ID of the user making the request.
            request_name: A name/title for the request.
            request: The request data/payload.
            request_input_schema: Optional JSON schema for validating the request.

        Returns:
            The created request object.

        Example:
            ```python
            req = client.create_request(
                tenant_id=1,
                requestor_id=5,
                request_name="Generate Audio",
                request={
                    "text": "Hello world",
                    "voice": "alloy",
                    "format": "mp3"
                }
            )
            ```
        """
        params = {
            "tenant_id": tenant_id,
            "requestor_id": requestor_id,
            "request_name": request_name,
            "request": request,
        }
        if request_input_schema is not None:
            params["request_input_schema"] = request_input_schema

        response = self._make_request("create_request", **params)
        return response["request"]

    def get_request(self, request_id: int) -> RequestModel:
        """
        Retrieve a request by ID.

        Args:
            request_id: The ID of the request to retrieve.

        Returns:
            The request object.

        Raises:
            NotFoundError: If the request is not found.
        """
        response = self._make_request("get_request", request_id=request_id)
        return response["request"]

    def list_requests(
        self,
        tenant_id: int,
        status: str | None = None,
        limit: int = 50,
        offset: int = 0,
    ) -> dict:
        """
        List requests for a tenant.

        Args:
            tenant_id: The ID of the tenant.
            status: Filter by request status (e.g., 'pending', 'completed').
            limit: Maximum number of requests to return.
            offset: Number of requests to skip for pagination.

        Returns:
            A dict containing 'requests' list and pagination info.
        """
        params = {
            "tenant_id": tenant_id,
            "limit": limit,
            "offset": offset,
        }
        if status is not None:
            params["status"] = status

        return self._make_request("list_requests", **params)

    # =========================================================================
    # API Key Operations
    # =========================================================================

    def list_api_keys(self, include_inactive: bool = False) -> list[APIKey]:
        """
        List API keys for your tenant.

        Note: For security, only the key prefix is returned, not the full key.

        Args:
            include_inactive: Whether to include revoked/inactive keys.

        Returns:
            A list of API key objects.
        """
        response = self._make_request(
            "list_api_keys",
            include_inactive=include_inactive,
        )
        return response["api_keys"]

    def revoke_api_key(self, api_key_id: int) -> dict:
        """
        Revoke an API key.

        Args:
            api_key_id: The ID of the API key to revoke.

        Returns:
            Confirmation of revocation.

        Raises:
            NotFoundError: If the API key is not found.
        """
        return self._make_request("revoke_api_key", api_key_id=api_key_id)

    # =========================================================================
    # Session Operations
    # =========================================================================

    def init(self) -> Session:
        """
        Initialize a session and validate authentication.

        This method validates the API key and returns information about
        the authenticated user and tenant. Use this to verify credentials
        and retrieve user/tenant context at the start of a session.

        Returns:
            A dict containing:
            - status: "ok" if successful
            - user: User information (id, email, name, etc.)
            - tenant: Tenant information (id, name, etc.)

        Raises:
            AuthenticationError: If the API key is invalid or expired.

        Example:
            ```python
            client = AshrLabsClient(api_key="tp_...", base_url="https://...")

            # Validate credentials and get user info
            session = client.init()
            print(f"Logged in as: {session['user']['email']}")
            print(f"Tenant: {session['tenant']['name']}")
            ```
        """
        return self._make_request("init")

    # =========================================================================
    # Utility Methods
    # =========================================================================

    def health_check(self) -> dict:
        """
        Check if the API is reachable and responding.

        Returns:
            A dict with status information.
        """
        return self._make_request("keep_alive")

    def __repr__(self) -> str:
        return f"AshrLabsClient(base_url='{self.base_url}', api_key='{self.api_key[:8]}...')"
