"""
Ashr Labs API Client.

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

import json
import time
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


_DEFAULT_BASE_URL = "https://api.ashr.io/testing-platform-api"


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

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

    The client automatically resolves your ``tenant_id`` and ``user_id``
    from the API key on first use, so you never need to pass them manually.

    Example:
        ```python
        from ashr_labs import AshrLabsClient

        client = AshrLabsClient(api_key="tp_your_api_key_here")

        # No tenant_id needed — auto-resolved from your API key
        datasets = client.list_datasets()

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

    def __init__(
        self,
        api_key: str,
        base_url: str = _DEFAULT_BASE_URL,
        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 (default: production URL).
            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

        # Lazy-loaded session context (populated on first call that needs it)
        self._tenant_id: int | None = None
        self._user_id: int | None = None
        self._session: dict | None = None

    @classmethod
    def from_env(cls, timeout: int = 30) -> "AshrLabsClient":
        """
        Create a client from environment variables.

        Reads ``ASHR_LABS_API_KEY`` (required) and ``ASHR_LABS_BASE_URL``
        (optional, defaults to production) from the environment.

        Returns:
            An authenticated AshrLabsClient.

        Raises:
            RuntimeError: If ASHR_LABS_API_KEY is not set.
        """
        import os

        api_key = os.environ.get("ASHR_LABS_API_KEY")
        if not api_key:
            raise RuntimeError(
                "ASHR_LABS_API_KEY environment variable is not set. "
                "Create an API key at https://app.ashr.io → API Keys."
            )
        base_url = os.environ.get("ASHR_LABS_BASE_URL", _DEFAULT_BASE_URL)
        return cls(api_key=api_key, base_url=base_url, timeout=timeout)

    def _ensure_session(self) -> None:
        """Lazily call init() and cache tenant_id / user_id."""
        if self._session is not None:
            return
        self._session = self._make_request("init")
        self._tenant_id = self._session["tenant"]["id"]
        self._user_id = self._session["user"]["id"]

    def _resolve_tenant_id(self, tenant_id: int | None) -> int:
        """Return the explicit tenant_id, or auto-resolve from session."""
        if tenant_id is not None:
            return tenant_id
        self._ensure_session()
        return self._tenant_id  # type: ignore[return-value]

    def _resolve_user_id(self, user_id: int | None) -> int:
        """Return the explicit user_id, or auto-resolve from session."""
        if user_id is not None:
            return user_id
        self._ensure_session()
        return self._user_id  # type: ignore[return-value]

    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 | None = None,
        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 (auto-resolved if omitted).
            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=self._resolve_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,
        dataset_id: int,
        result: dict[str, Any],
        tenant_id: int | None = None,
        runner_id: int | None = None,
    ) -> Run:
        """
        Create a new test run.

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

        Returns:
            The created run object.

        Example:
            ```python
            run = client.create_run(
                dataset_id=42,
                result={
                    "score": 0.95,
                    "status": "passed",
                    "metrics": {"accuracy": 0.98, "latency_ms": 150}
                }
            )
            ```
        """
        params: dict[str, Any] = {
            "tenant_id": self._resolve_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,
        dataset_id: int | None = None,
        tenant_id: int | None = None,
        limit: int = 50,
        offset: int = 0,
    ) -> dict:
        """
        List runs for a tenant or dataset.

        Args:
            dataset_id: Filter by dataset ID.
            tenant_id: Filter by tenant ID (auto-resolved if omitted).
            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: dict[str, Any] = {
            "tenant_id": self._resolve_tenant_id(tenant_id),
            "limit": limit,
            "offset": offset,
        }
        if dataset_id is not None:
            params["dataset_id"] = dataset_id

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

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

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

        The ``request`` dict contains the generation config that describes
        your agent and the test scenarios to generate.  The
        ``request_input_schema`` is a JSON Schema that the backend uses to
        validate that config; if you include agent tools, put them in
        ``request_input_schema["tools"]`` so they're auto-saved as skill
        templates.

        Args:
            request_name: A name/title for the request.
            request: The generation config (see below).
            request_input_schema: JSON Schema for validating ``request``.
                If omitted a permissive default schema is sent.
            tenant_id: The ID of the tenant (auto-resolved if omitted).
            requestor_id: The ID of the user making the request (auto-resolved if omitted).

        Returns:
            The created request object.

        **Generation config structure** (the ``request`` dict)::

            {
                "metadata": {"dataset_name": str, "description": str},
                "agent": {
                    "name": str,
                    "description": str,
                    "system_prompt": str,
                    "tools": [{"name": str, "description": str, "parameters": {...}}, ...],
                    "accepted_inputs": {"text": bool, "audio": bool, "file": bool, ...},
                    "output_format": {"type": "text"},
                },
                "context": {
                    "domain": str,
                    "use_case": str,
                    "scenario_context": str,
                },
                "test_config": {
                    "num_variations": int,
                    "coverage": {"happy_path": bool, "edge_cases": bool, ...},
                },
                "generation_options": {
                    "generate_audio": bool,
                    "generate_files": bool,
                    "generate_simulations": bool,
                },
            }

        Example:
            ```python
            req = client.create_request(
                request_name="Support Agent Eval",
                request={
                    "metadata": {"dataset_name": "Support Eval"},
                    "agent": {
                        "name": "Support Bot",
                        "description": "Answers customer questions",
                        "system_prompt": "You are a helpful support agent.",
                        "tools": [
                            {"name": "lookup_order", "description": "Look up an order",
                             "parameters": {"type": "object", "properties": {"order_id": {"type": "string"}}, "required": ["order_id"]}},
                        ],
                        "accepted_inputs": {"text": True, "audio": False, "file": False, "image": False, "video": False},
                        "output_format": {"type": "text"},
                    },
                    "context": {
                        "domain": "ecommerce",
                        "use_case": "Customers asking about orders",
                        "scenario_context": "An online retail store",
                    },
                    "test_config": {"num_variations": 3, "coverage": {"happy_path": True, "edge_cases": True}},
                    "generation_options": {"generate_audio": False, "generate_files": False, "generate_simulations": False},
                },
                request_input_schema={
                    "type": "object",
                    "properties": {"agent": {"type": "object"}, "context": {"type": "object"}},
                    "required": ["agent", "context"],
                    "tools": [{"name": "lookup_order", "description": "Look up an order",
                               "parameters": {"type": "object", "properties": {"order_id": {"type": "string"}}}}],
                },
            )
            ```
        """
        # The backend requires request_input_schema; send a permissive
        # default so callers aren't forced to build one manually.
        if request_input_schema is None:
            request_input_schema = {
                "type": "object",
                "properties": {},
            }
            # Auto-populate tools from the request's agent config
            tools = (request.get("agent") or {}).get("tools")
            if tools:
                request_input_schema["tools"] = tools

        params: dict[str, Any] = {
            "tenant_id": self._resolve_tenant_id(tenant_id),
            "requestor_id": self._resolve_user_id(requestor_id),
            "request_name": request_name,
            "request": request,
            "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 | None = None,
        status: str | None = None,
        limit: int = 50,
        offset: int = 0,
    ) -> dict:
        """
        List requests for a tenant.

        Args:
            tenant_id: The ID of the tenant (auto-resolved if omitted).
            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: dict[str, Any] = {
            "tenant_id": self._resolve_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, tenant_name, etc.)

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

        Example:
            ```python
            client = AshrLabsClient(api_key="tp_...")

            # Validate credentials and get user info
            session = client.init()
            print(f"Logged in as: {session['user']['email']}")
            print(f"Tenant: {session['tenant']['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")

    # =========================================================================
    # Convenience / Workflow Methods
    # =========================================================================

    def wait_for_request(
        self,
        request_id: int,
        timeout: int = 600,
        poll_interval: int = 5,
    ) -> RequestModel:
        """Block until a request reaches a terminal state (completed/failed).

        Args:
            request_id: The ID of the request to poll.
            timeout: Maximum seconds to wait before raising TimeoutError.
            poll_interval: Seconds between polls.

        Returns:
            The final request object.

        Raises:
            TimeoutError: If the request doesn't finish within ``timeout`` seconds.
            AshrLabsError: If the request fails.
        """
        start = time.time()
        while time.time() - start < timeout:
            req = self.get_request(request_id=request_id)
            status = req.get("request_status", "")
            if status == "completed":
                return req
            if status == "failed":
                error_msg = req.get("error", "Request failed")
                raise AshrLabsError(f"Request {request_id} failed: {error_msg}")
            time.sleep(poll_interval)
        raise TimeoutError(
            f"Request {request_id} did not complete within {timeout}s"
        )

    def generate_dataset(
        self,
        request_name: str,
        config: dict[str, Any],
        request_input_schema: dict[str, Any] | None = None,
        timeout: int = 600,
        poll_interval: int = 5,
    ) -> tuple[int, dict]:
        """Create a dataset generation request, wait for it, and fetch the result.

        This is a high-level convenience method that combines
        ``create_request`` + ``wait_for_request`` + ``get_dataset``.

        Args:
            request_name: A name/title for the request.
            config: The generation config dict (same as ``create_request``'s
                ``request`` parameter).
            request_input_schema: Optional JSON Schema for validation.
            timeout: Maximum seconds to wait for generation.
            poll_interval: Seconds between polls.

        Returns:
            A tuple of ``(dataset_id, dataset_source)`` where
            ``dataset_source`` is the dict containing ``"runs"``.

        Raises:
            TimeoutError: If generation doesn't finish in time.
            AshrLabsError: If generation fails.
        """
        req = self.create_request(
            request_name=request_name,
            request=config,
            request_input_schema=request_input_schema,
        )
        request_id = req["id"]

        self.wait_for_request(
            request_id, timeout=timeout, poll_interval=poll_interval
        )

        # Fetch the most recent dataset (the one just generated)
        resp = self.list_datasets(limit=1)
        datasets = resp.get("datasets", [])
        if not datasets:
            raise AshrLabsError("No datasets found after generation completed")

        dataset_id = datasets[0]["id"]
        full_ds = self.get_dataset(dataset_id=dataset_id, include_signed_urls=False)
        source = full_ds.get("dataset_source", {})
        return dataset_id, source

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