"""Enhanced screenshot service client with polling and retry logic.

Wraps the auto-generated screenshot-client-core and adds:
- Polling with configurable timeout/interval
- Automatic retry on transient errors
- Environment-based configuration
- Domain-specific exceptions
"""

from __future__ import annotations

import os
from collections.abc import Callable
from typing import Any

import httpx
from infra_core.polling import PollingConfig, StatusPoller
from screenshot_client_core import AuthenticatedClient, Client
from screenshot_client_core.api.batches import list_screenshot_runs, start_screenshot_batch
from screenshot_client_core.api.results import get_screenshot_result
from screenshot_client_core.api.runs import cancel_screenshot_run, get_screenshot_run
from screenshot_client_core.models import (
    RunListResponse,
    ScreenshotBatchRequest,
    ScreenshotResultManifest,
    ScreenshotRunHandle,
    ScreenshotRunStatus,
)
from screenshot_client_core.types import UNSET, Unset

from screenshot_batch.exceptions import (
    ScreenshotAPIError,
    ScreenshotConfigError,
)

_DEFAULT_TIMEOUT = 30.0
_DEFAULT_POLL_TIMEOUT = 600.0
_DEFAULT_POLL_INTERVAL = 5.0


class ScreenshotServiceClient:
    """Enhanced client with polling, retry, and convenience methods.

    This wraps the auto-generated screenshot-client-core and adds:
    - Polling with configurable timeout/interval
    - Automatic retry on transient errors
    - Environment-based configuration
    - Domain-specific exceptions

    Args:
        base_url: Base URL of the screenshot service.
        api_key: Optional API key for authentication.
        timeout: Default timeout for HTTP requests in seconds. Defaults to 30.0.

    Example:
        >>> client = ScreenshotServiceClient.from_env()
        >>> handle = client.start_batch("my-batch", jobs=[...])
        >>> result = client.wait_for_completion(handle, timeout=600)
    """

    def __init__(
        self,
        base_url: str,
        *,
        api_key: str | None = None,
        timeout: float = _DEFAULT_TIMEOUT,
    ):
        # Use Client (no auth) when no api_key, AuthenticatedClient otherwise
        # AuthenticatedClient with empty token causes "Illegal header value b'Bearer '"
        core_client: Client | AuthenticatedClient
        if api_key:
            core_client = AuthenticatedClient(
                base_url=base_url,
                token=api_key,
                timeout=httpx.Timeout(timeout),
            ).with_headers({"x-functions-key": api_key})
        else:
            core_client = Client(
                base_url=base_url,
                timeout=httpx.Timeout(timeout),
            )
        self._core_client = core_client

        self.base_url = base_url
        self.timeout = timeout

    @classmethod
    def from_env(
        cls,
        *,
        url_env: str = "SCREENSHOT_FUNC_BASE_URL",
        api_key_env: str = "SCREENSHOT_FUNC_API_KEY",
        timeout: float = _DEFAULT_TIMEOUT,
    ) -> ScreenshotServiceClient:
        """Create client from environment variables.

        Args:
            url_env: Environment variable name for service URL. Defaults to SCREENSHOT_FUNC_BASE_URL.
            api_key_env: Environment variable name for API key. Defaults to SCREENSHOT_FUNC_API_KEY.
            timeout: Default timeout in seconds. Defaults to 30.0.

        Returns:
            Configured client instance.

        Raises:
            ScreenshotConfigError: If required environment variables are missing.

        Example:
            >>> client = ScreenshotServiceClient.from_env()
        """
        base_url = os.getenv(url_env)
        if not base_url:
            raise ScreenshotConfigError(f"Environment variable {url_env} not set")

        api_key = os.getenv(api_key_env)  # Optional

        return cls(base_url=base_url, api_key=api_key, timeout=timeout)

    def start_batch(
        self,
        batch_id: str,
        payload: dict[str, Any],
    ) -> ScreenshotRunHandle:
        """Start a screenshot batch run.

        Submit a batch of screenshot jobs for asynchronous processing.

        Args:
            batch_id: Unique identifier for this batch.
            payload: Request payload with 'jobs' list and optional 'concurrency'.

        Returns:
            Handle with batch_id, job_id, and status URL.

        Raises:
            HTTPError: If batch submission fails.

        Example:
            >>> handle = client.start_batch(
            ...     "my-batch",
            ...     {"jobs": [{"job_id": "1", "url": "https://example.com"}]}
            ... )
        """
        request = ScreenshotBatchRequest.from_dict(payload)
        response = start_screenshot_batch.sync_detailed(
            batch_id=batch_id,
            client=self._core_client,  # type: ignore[arg-type]  # type: ignore[arg-type]
            body=request,
        )

        if response.status_code != 202:
            raise ScreenshotAPIError(
                status_code=response.status_code,
                content=response.content,
                message=f"Failed to start batch: unexpected status {response.status_code}",
            )

        if not isinstance(response.parsed, ScreenshotRunHandle):
            raise ScreenshotAPIError(
                status_code=response.status_code,
                content=response.content,
                message=f"Invalid response type: {type(response.parsed)}",
            )

        return response.parsed

    def get_run(
        self,
        batch_id: str,
        job_id: str,
        *,
        pending_on_404: bool = False,
    ) -> dict[str, Any]:
        """Get current status of a batch run.

        Args:
            batch_id: Batch identifier.
            job_id: Job identifier.
            pending_on_404: If True, return synthetic pending status on 404
                instead of raising. Useful for polling after immediate submission.

        Returns:
            Status dictionary with fields: status, job_total, job_completed, etc.

        Raises:
            ScreenshotAPIError: If status fetch fails (unless pending_on_404=True and 404).

        Example:
            >>> status = client.get_run("my-batch", "550e8400-e29b-...")
            >>> print(status.status)
            running
        """
        response = get_screenshot_run.sync_detailed(
            batch_id=batch_id,
            job_id=job_id,
            client=self._core_client,  # type: ignore[arg-type]
        )

        # Handle 404 gracefully for polling after immediate submission
        # The durable orchestration may not have created the run record yet
        if response.status_code == 404 and pending_on_404:
            return {
                "batch_id": batch_id,
                "job_id": job_id,
                "status": "pending",
            }

        if response.status_code != 200:
            raise ScreenshotAPIError(
                status_code=response.status_code,
                content=response.content,
                message=f"Failed to get run status: unexpected status {response.status_code}",
            )

        if not isinstance(response.parsed, ScreenshotRunStatus):
            raise ScreenshotAPIError(
                status_code=response.status_code,
                content=response.content,
                message=f"Invalid response type: {type(response.parsed)}",
            )

        result: dict[str, Any] = response.parsed.to_dict()
        return result

    def list_runs(
        self,
        batch_id: str,
        *,
        limit: int = 50,
        continuation: str | None = None,
    ) -> list[dict[str, Any]]:
        """List all runs for a batch.

        Args:
            batch_id: Batch identifier.
            limit: Maximum number of items to return. Defaults to 50.
            continuation: Continuation token from previous response.

        Returns:
            List of run status dictionaries.

        Example:
            >>> runs = client.list_runs("my-batch", limit=100)
        """
        response = list_screenshot_runs.sync_detailed(
            batch_id=batch_id,
            client=self._core_client,  # type: ignore[arg-type]
            limit=limit,
            continuation=continuation if continuation is not None else UNSET,
        )

        if response.status_code != 200:
            raise ScreenshotAPIError(
                status_code=response.status_code,
                content=response.content,
                message=f"Failed to list runs: unexpected status {response.status_code}",
            )

        # Extract items from response
        parsed = response.parsed
        if isinstance(parsed, RunListResponse):
            items = parsed.items
            if not isinstance(items, Unset):
                return [item.to_dict() for item in items]
        return []

    def get_result(self, batch_id: str, job_id: str) -> dict[str, Any]:
        """Fetch result manifest for a completed run.

        Args:
            batch_id: Batch identifier.
            job_id: Job identifier.

        Returns:
            Result manifest with storage paths and file listings.

        Raises:
            HTTPError: If result fetch fails.

        Example:
            >>> result = client.get_result("my-batch", "550e8400-e29b-...")
            >>> print(result["manifest"]["storage_path"])
        """
        response = get_screenshot_result.sync_detailed(
            batch_id=batch_id,
            job_id=job_id,
            client=self._core_client,  # type: ignore[arg-type]
        )

        if response.status_code != 200:
            raise ScreenshotAPIError(
                status_code=response.status_code,
                content=response.content,
                message=f"Failed to get result: unexpected status {response.status_code}",
            )

        if not isinstance(response.parsed, ScreenshotResultManifest):
            raise ScreenshotAPIError(
                status_code=response.status_code,
                content=response.content,
                message=f"Invalid response type: {type(response.parsed)}",
            )

        result: dict[str, Any] = response.parsed.to_dict()
        return result

    def cancel_run(self, batch_id: str, job_id: str) -> None:
        """Cancel an in-flight run.

        Args:
            batch_id: Batch identifier.
            job_id: Job identifier.

        Raises:
            HTTPError: If cancellation fails.

        Example:
            >>> client.cancel_run("my-batch", "550e8400-e29b-...")
        """
        response = cancel_screenshot_run.sync_detailed(
            batch_id=batch_id,
            job_id=job_id,
            client=self._core_client,  # type: ignore[arg-type]
        )

        if response.status_code != 200:
            raise ScreenshotAPIError(
                status_code=response.status_code,
                content=response.content,
                message=f"Failed to cancel run: unexpected status {response.status_code}",
            )

    def poll_run(
        self,
        batch_id: str,
        job_id: str,
        *,
        timeout: float = _DEFAULT_POLL_TIMEOUT,
        poll_interval: float = _DEFAULT_POLL_INTERVAL,
        on_poll: Callable[[dict[str, Any]], None] | None = None,
    ) -> dict[str, Any]:
        """Poll run status until completion or timeout.

        Uses the reusable polling utilities from infra_core.polling.

        Args:
            batch_id: Batch identifier.
            job_id: Job identifier.
            timeout: Maximum time to wait in seconds. Defaults to 600.0.
            poll_interval: Time between status checks in seconds. Defaults to 5.0.
            on_poll: Optional callback(status_dict) called on each poll.

        Returns:
            Final status dictionary when terminal state is reached.

        Raises:
            PollingTimeoutError: If timeout is reached before completion.
            PollingFailureError: If batch fails.

        Example:
            >>> result = client.poll_run(
            ...     "my-batch",
            ...     "550e8400-e29b-...",
            ...     timeout=600,
            ...     on_poll=lambda s: print(f"Progress: {s.get('job_completed')}/{s.get('job_total')}")
            ... )
        """
        config = PollingConfig(
            timeout=timeout,
            poll_interval=poll_interval,
            terminal_statuses=frozenset({"completed", "failed", "cancelled"}),
            failure_statuses=frozenset({"failed"}),
        )

        def fetch_status() -> dict[str, Any]:
            # Use pending_on_404=True to handle race condition where the
            # durable orchestration hasn't created the run record yet
            return self.get_run(batch_id, job_id, pending_on_404=True)

        poller = StatusPoller(
            fetch_fn=fetch_status,
            config=config,
        )

        return poller.poll(on_poll=on_poll)
