"""Sandai Operator client implementation."""

import asyncio
import logging
try:
    from functools import cached_property
except ImportError:  # pragma: no cover
    from cached_property import cached_property
from typing import Any, Dict, List, Optional, Tuple

from celery.result import AsyncResult

from .config import (
    VALID_PRIORITIES,
    CeleryConfig,
    ClientConfig,
    get_celery_config,
    get_default_celery_config,
    resolve_client_config,
)
from .exceptions import InvalidPriorityError, OperatorError
from .protocol import (
    build_task_payload,
    inspect_probe_result,
    parse_desc_result,
    parse_task_result,
    raise_unexpected_task_error,
)
from .transport import CeleryTransport
from .async_transport import AsyncTransportAdapter
from .types import (
    AsyncSubmittedTask,
    AsyncTaskHandle,
    BatchTaskResult,
    HealthCheckResult,
    OperatorDescription,
    SubmittedTask,
)

logger = logging.getLogger(__name__)


LegacyTaskResult = Tuple[List[Dict[str, Any]], Dict[str, Any]]
LegacyBatchTasks = List[Tuple[List[Dict[str, Any]], Dict[str, Any]]]


class OperatorClient:
    """Client for sending tasks to a Sandai operator over Celery."""

    def __init__(
        self,
        operator_name: str,
        operator_version: str,
        broker_url: Optional[str] = None,
        result_backend: Optional[str] = None,
        timeout: int = 28800,
        debug: bool = False,
        priority: str = "normal",
        cancel_key: Optional[str] = None,
        auto_forget_result: Optional[bool] = None,
    ):
        self.valid_priorities = list(VALID_PRIORITIES)
        self._validate_priority(priority)

        self.config = resolve_client_config(
            operator_name,
            operator_version,
            broker_url=broker_url,
            result_backend=result_backend,
            timeout=timeout,
            debug=debug,
            priority=priority,
            cancel_key=cancel_key,
            auto_forget_result=auto_forget_result,
        )
        self.transport = CeleryTransport(self.config, logger=logger)
        self.celery_app = self.transport.app

        self.operator_name = self.config.operator_name
        self.operator_version = self.config.operator_version
        self.timeout = self.config.timeout
        self.debug = self.config.debug
        self.priority = self.config.priority
        self.cancel_key = self.config.cancel_key
        self.auto_forget_result = self.config.auto_forget_result
        self.task_name = self.config.task_name
        self.queue_names = self.config.queue_names
        self.broker_url = self.config.broker_url
        self.result_backend = self.config.result_backend

        if self.debug:
            logger.debug(
                "OperatorClient initialized operator=%s version=%s broker=%s backend=%s queues=%s",
                self.operator_name,
                self.operator_version,
                self.broker_url,
                self.result_backend,
                list(self.queue_names.values()),
            )

    def _validate_priority(self, priority: str) -> None:
        if priority not in self.valid_priorities:
            raise InvalidPriorityError(priority, self.valid_priorities)

    def _effective_timeout(self, timeout: Optional[int]) -> int:
        return timeout if timeout is not None else self.timeout

    def _forget_result_if_needed(self, async_result: Optional[AsyncResult]) -> None:
        if not self.auto_forget_result or async_result is None:
            return

        try:
            self.transport.forget(async_result)
            if self.debug:
                logger.debug("forgot result key task_id=%s", async_result.id)
        except Exception as forget_error:  # pragma: no cover - defensive logging only
            if self.debug:
                logger.debug(
                    "failed to forget result key task_id=%s error=%s",
                    getattr(async_result, "id", None),
                    forget_error,
                )

    def _submit_payload(self, payload: Dict[str, Any], *, priority: str) -> AsyncResult:
        queue_name = self.queue_names[priority]
        return self.transport.send_task(self.task_name, kwargs=payload, queue=queue_name)

    def _build_task_payload(
        self,
        input_artifacts: List[Dict[str, Any]],
        options: Dict[str, Any],
        *,
        timeout: Optional[int],
        cancel_key: Optional[str],
        persistent_output: bool,
    ) -> Dict[str, Any]:
        return build_task_payload(
            input_artifacts,
            options,
            timeout=timeout,
            default_timeout=self.timeout,
            cancel_key=cancel_key,
            default_cancel_key=self.cancel_key,
            persistent_output=persistent_output,
        )

    def _send_task_async(
        self,
        input_artifacts: List[Dict[str, Any]],
        options: Dict[str, Any],
        priority: Optional[str] = None,
        *,
        timeout: Optional[int] = None,
        cancel_key: Optional[str] = None,
        persistent_output: bool = False,
    ) -> AsyncResult:
        resolved_priority = priority or self.priority
        self._validate_priority(resolved_priority)
        payload = self._build_task_payload(
            input_artifacts,
            options,
            timeout=timeout,
            cancel_key=cancel_key,
            persistent_output=persistent_output,
        )
        return self._submit_payload(payload, priority=resolved_priority)

    def _wait_for_task_result(
        self, async_result: AsyncResult, *, timeout: int
    ) -> LegacyTaskResult:
        try:
            raw_result = self.transport.get_result(async_result, timeout=timeout)
            return parse_task_result(async_result.id, raw_result).to_legacy_tuple()
        except OperatorError:
            raise
        except Exception as exc:
            raise_unexpected_task_error(async_result.id, timeout, exc)

    def _wait_for_desc(self, async_result: AsyncResult, *, timeout: int) -> OperatorDescription:
        try:
            raw_result = self.transport.get_result(async_result, timeout=timeout)
            return parse_desc_result(async_result.id, raw_result)
        except OperatorError:
            raise
        except Exception as exc:
            raise_unexpected_task_error(async_result.id, timeout, exc)

    @cached_property
    def desc(self) -> Dict[str, Any]:
        async_result = self._submit_payload({"task_type": "desc"}, priority="high")
        try:
            return self._wait_for_desc(async_result, timeout=self.timeout).to_dict()
        finally:
            self._forget_result_if_needed(async_result)

    def async_call(
        self,
        input_artifacts: List[Dict[str, Any]],
        options: Dict[str, Any],
        priority: Optional[str] = None,
        timeout: Optional[int] = None,
        cancel_key: Optional[str] = None,
        persistent_output: bool = False,
    ) -> AsyncResult:
        return self._send_task_async(
            input_artifacts,
            options,
            priority=priority,
            timeout=timeout,
            cancel_key=cancel_key,
            persistent_output=persistent_output,
        )

    def sync(
        self,
        input_artifacts: List[Dict[str, Any]],
        options: Dict[str, Any],
        priority: Optional[str] = None,
        *,
        timeout: Optional[int] = None,
        cancel_key: Optional[str] = None,
        persistent_output: bool = False,
    ) -> LegacyTaskResult:
        async_result = self._send_task_async(
            input_artifacts,
            options,
            priority=priority,
            timeout=timeout,
            cancel_key=cancel_key,
            persistent_output=persistent_output,
        )
        effective_timeout = self._effective_timeout(timeout)

        try:
            return self._wait_for_task_result(async_result, timeout=effective_timeout)
        finally:
            self._forget_result_if_needed(async_result)

    def batch_submit(
        self,
        tasks: LegacyBatchTasks,
        *,
        priority: Optional[str] = None,
        timeout: Optional[int] = None,
        cancel_key: Optional[str] = None,
        persistent_output: bool = False,
    ) -> List[SubmittedTask]:
        submissions: List[SubmittedTask] = []
        for index, (input_artifacts, options) in enumerate(tasks):
            try:
                async_result = self._send_task_async(
                    input_artifacts,
                    options,
                    priority=priority,
                    timeout=timeout,
                    cancel_key=cancel_key,
                    persistent_output=persistent_output,
                )
                submissions.append(SubmittedTask(index=index, async_result=async_result))
            except Exception as exc:
                submissions.append(SubmittedTask(index=index, error=exc))
        return submissions

    def batch_async(
        self,
        tasks: LegacyBatchTasks,
        *,
        priority: Optional[str] = None,
        timeout: Optional[int] = None,
        cancel_key: Optional[str] = None,
        persistent_output: bool = False,
    ) -> List[Optional[AsyncResult]]:
        """Legacy best-effort batch submit API kept for compatibility."""

        submissions = self.batch_submit(
            tasks,
            priority=priority,
            timeout=timeout,
            cancel_key=cancel_key,
            persistent_output=persistent_output,
        )
        return [submission.async_result for submission in submissions]

    def batch_collect(
        self, submissions: List[SubmittedTask], *, timeout: Optional[int] = None
    ) -> List[BatchTaskResult]:
        effective_timeout = self._effective_timeout(timeout)
        results: List[BatchTaskResult] = []

        for submission in submissions:
            if not submission.submitted or submission.async_result is None:
                results.append(
                    BatchTaskResult(
                        index=submission.index,
                        task_id=submission.task_id,
                        submitted=False,
                        error=submission.error,
                    )
                )
                continue

            async_result = submission.async_result
            try:
                artifacts, task_results = self._wait_for_task_result(
                    async_result,
                    timeout=effective_timeout,
                )
                results.append(
                    BatchTaskResult(
                        index=submission.index,
                        task_id=submission.task_id,
                        submitted=True,
                        artifacts=artifacts,
                        results=task_results,
                    )
                )
            except Exception as exc:
                results.append(
                    BatchTaskResult(
                        index=submission.index,
                        task_id=submission.task_id,
                        submitted=True,
                        error=exc,
                    )
                )
            finally:
                self._forget_result_if_needed(async_result)

        return results

    def batch_sync_detailed(
        self,
        tasks: LegacyBatchTasks,
        *,
        priority: Optional[str] = None,
        timeout: Optional[int] = None,
        cancel_key: Optional[str] = None,
        persistent_output: bool = False,
    ) -> List[BatchTaskResult]:
        submissions = self.batch_submit(
            tasks,
            priority=priority,
            timeout=timeout,
            cancel_key=cancel_key,
            persistent_output=persistent_output,
        )
        return self.batch_collect(submissions, timeout=timeout)

    def batch_sync(
        self,
        tasks: LegacyBatchTasks,
        *,
        priority: Optional[str] = None,
        timeout: Optional[int] = None,
        cancel_key: Optional[str] = None,
        persistent_output: bool = False,
    ) -> List[LegacyTaskResult]:
        """Legacy best-effort batch sync API kept for compatibility."""

        detailed_results = self.batch_sync_detailed(
            tasks,
            priority=priority,
            timeout=timeout,
            cancel_key=cancel_key,
            persistent_output=persistent_output,
        )
        return [result.to_legacy_tuple() for result in detailed_results]

    def get_queue_info(self) -> Dict[str, str]:
        return self.queue_names.copy()

    def probe(self, timeout: int = 10) -> HealthCheckResult:
        try:
            async_result = self._send_task_async(
                input_artifacts=[],
                options={"test": True},
                priority="low",
            )
        except Exception as exc:
            return HealthCheckResult(
                submit_ok=False,
                retrieval_ok=False,
                response_ok=False,
                error=str(exc),
            )

        try:
            raw_result = self.transport.get_result(async_result, timeout=timeout)
            response_ok, probe_error = inspect_probe_result(raw_result)
            return HealthCheckResult(
                submit_ok=True,
                retrieval_ok=True,
                response_ok=response_ok,
                task_id=async_result.id,
                error=probe_error,
            )
        except Exception as exc:
            return HealthCheckResult(
                submit_ok=True,
                retrieval_ok=False,
                response_ok=False,
                task_id=async_result.id,
                error=str(exc),
            )
        finally:
            self._forget_result_if_needed(async_result)

    def ping(self, timeout: int = 10) -> bool:
        """Compatibility helper that confirms the submit path is reachable."""

        return self.probe(timeout=timeout).submit_ok


class AsyncOperatorClient:
    """Async client for sending tasks to a Sandai operator."""

    def __init__(
        self,
        operator_name: str,
        operator_version: str,
        broker_url: Optional[str] = None,
        result_backend: Optional[str] = None,
        timeout: int = 28800,
        debug: bool = False,
        priority: str = "normal",
        cancel_key: Optional[str] = None,
        auto_forget_result: Optional[bool] = None,
    ):
        self.valid_priorities = list(VALID_PRIORITIES)
        self._validate_priority(priority)

        self.config = resolve_client_config(
            operator_name,
            operator_version,
            broker_url=broker_url,
            result_backend=result_backend,
            timeout=timeout,
            debug=debug,
            priority=priority,
            cancel_key=cancel_key,
            auto_forget_result=auto_forget_result,
        )
        self.transport = AsyncTransportAdapter(self.config, logger_=logger)
        self.celery_app = self.transport.app

        self.operator_name = self.config.operator_name
        self.operator_version = self.config.operator_version
        self.timeout = self.config.timeout
        self.debug = self.config.debug
        self.priority = self.config.priority
        self.cancel_key = self.config.cancel_key
        self.auto_forget_result = self.config.auto_forget_result
        self.task_name = self.config.task_name
        self.queue_names = self.config.queue_names
        self.broker_url = self.config.broker_url
        self.result_backend = self.config.result_backend
        self._desc_cache: Optional[Dict[str, Any]] = None
        self._desc_lock = asyncio.Lock()

    def _validate_priority(self, priority: str) -> None:
        if priority not in self.valid_priorities:
            raise InvalidPriorityError(priority, self.valid_priorities)

    def _effective_timeout(self, timeout: Optional[int]) -> int:
        return timeout if timeout is not None else self.timeout

    async def aclose(self) -> None:
        await self.transport.aclose()

    async def __aenter__(self) -> "AsyncOperatorClient":
        return self

    async def __aexit__(self, exc_type, exc, tb) -> None:
        await self.aclose()

    async def _forget_result_if_needed(self, handle: Optional[AsyncTaskHandle]) -> None:
        if not self.auto_forget_result or handle is None:
            return

        try:
            await handle.forget()
            if self.debug:
                logger.debug("forgot result key task_id=%s", handle.task_id)
        except Exception as forget_error:  # pragma: no cover
            if self.debug:
                logger.debug(
                    "failed to forget result key task_id=%s error=%s",
                    getattr(handle, "task_id", None),
                    forget_error,
                )

    def _build_task_payload(
        self,
        input_artifacts: List[Dict[str, Any]],
        options: Dict[str, Any],
        *,
        timeout: Optional[int],
        cancel_key: Optional[str],
        persistent_output: bool,
    ) -> Dict[str, Any]:
        return build_task_payload(
            input_artifacts,
            options,
            timeout=timeout,
            default_timeout=self.timeout,
            cancel_key=cancel_key,
            default_cancel_key=self.cancel_key,
            persistent_output=persistent_output,
        )

    async def _submit_payload(self, payload: Dict[str, Any], *, priority: str) -> AsyncTaskHandle:
        queue_name = self.queue_names[priority]
        async_result = await self.transport.send_task(self.task_name, kwargs=payload, queue=queue_name)
        return AsyncTaskHandle(async_result=async_result, transport=self.transport)

    async def _send_task_async(
        self,
        input_artifacts: List[Dict[str, Any]],
        options: Dict[str, Any],
        priority: Optional[str] = None,
        *,
        timeout: Optional[int] = None,
        cancel_key: Optional[str] = None,
        persistent_output: bool = False,
    ) -> AsyncTaskHandle:
        resolved_priority = priority or self.priority
        self._validate_priority(resolved_priority)
        payload = self._build_task_payload(
            input_artifacts,
            options,
            timeout=timeout,
            cancel_key=cancel_key,
            persistent_output=persistent_output,
        )
        return await self._submit_payload(payload, priority=resolved_priority)

    async def _wait_for_task_result(
        self, handle: AsyncTaskHandle, *, timeout: int
    ) -> LegacyTaskResult:
        try:
            raw_result = await self.transport.get_result(handle.async_result, timeout=timeout)
            return parse_task_result(handle.task_id or "", raw_result).to_legacy_tuple()
        except OperatorError:
            raise
        except Exception as exc:
            raise_unexpected_task_error(handle.task_id or "", timeout, exc)

    async def _wait_for_desc(self, handle: AsyncTaskHandle, *, timeout: int) -> OperatorDescription:
        try:
            raw_result = await self.transport.get_result(handle.async_result, timeout=timeout)
            return parse_desc_result(handle.task_id or "", raw_result)
        except OperatorError:
            raise
        except Exception as exc:
            raise_unexpected_task_error(handle.task_id or "", timeout, exc)

    async def desc(self) -> Dict[str, Any]:
        if self._desc_cache is not None:
            return self._desc_cache

        async with self._desc_lock:
            if self._desc_cache is not None:
                return self._desc_cache
            handle = await self._submit_payload({"task_type": "desc"}, priority="high")
            try:
                self._desc_cache = (await self._wait_for_desc(handle, timeout=self.timeout)).to_dict()
                return self._desc_cache
            finally:
                await self._forget_result_if_needed(handle)

    async def async_call(
        self,
        input_artifacts: List[Dict[str, Any]],
        options: Dict[str, Any],
        priority: Optional[str] = None,
        timeout: Optional[int] = None,
        cancel_key: Optional[str] = None,
        persistent_output: bool = False,
    ) -> AsyncTaskHandle:
        return await self._send_task_async(
            input_artifacts,
            options,
            priority=priority,
            timeout=timeout,
            cancel_key=cancel_key,
            persistent_output=persistent_output,
        )

    async def sync(
        self,
        input_artifacts: List[Dict[str, Any]],
        options: Dict[str, Any],
        priority: Optional[str] = None,
        *,
        timeout: Optional[int] = None,
        cancel_key: Optional[str] = None,
        persistent_output: bool = False,
    ) -> LegacyTaskResult:
        handle = await self._send_task_async(
            input_artifacts,
            options,
            priority=priority,
            timeout=timeout,
            cancel_key=cancel_key,
            persistent_output=persistent_output,
        )
        effective_timeout = self._effective_timeout(timeout)
        try:
            return await self._wait_for_task_result(handle, timeout=effective_timeout)
        finally:
            await self._forget_result_if_needed(handle)

    async def batch_submit(
        self,
        tasks: LegacyBatchTasks,
        *,
        priority: Optional[str] = None,
        timeout: Optional[int] = None,
        cancel_key: Optional[str] = None,
        persistent_output: bool = False,
    ) -> List[AsyncSubmittedTask]:
        submissions: List[AsyncSubmittedTask] = []
        for index, (input_artifacts, options) in enumerate(tasks):
            try:
                handle = await self._send_task_async(
                    input_artifacts,
                    options,
                    priority=priority,
                    timeout=timeout,
                    cancel_key=cancel_key,
                    persistent_output=persistent_output,
                )
                submissions.append(AsyncSubmittedTask(index=index, handle=handle))
            except Exception as exc:
                submissions.append(AsyncSubmittedTask(index=index, error=exc))
        return submissions

    async def batch_async(
        self,
        tasks: LegacyBatchTasks,
        *,
        priority: Optional[str] = None,
        timeout: Optional[int] = None,
        cancel_key: Optional[str] = None,
        persistent_output: bool = False,
    ) -> List[Optional[AsyncTaskHandle]]:
        submissions = await self.batch_submit(
            tasks,
            priority=priority,
            timeout=timeout,
            cancel_key=cancel_key,
            persistent_output=persistent_output,
        )
        return [submission.handle for submission in submissions]

    async def batch_collect(
        self, submissions: List[AsyncSubmittedTask], *, timeout: Optional[int] = None
    ) -> List[BatchTaskResult]:
        effective_timeout = self._effective_timeout(timeout)
        results: List[BatchTaskResult] = []

        for submission in submissions:
            if not submission.submitted or submission.handle is None:
                results.append(
                    BatchTaskResult(
                        index=submission.index,
                        task_id=submission.task_id,
                        submitted=False,
                        error=submission.error,
                    )
                )
                continue

            handle = submission.handle
            try:
                artifacts, task_results = await self._wait_for_task_result(
                    handle,
                    timeout=effective_timeout,
                )
                results.append(
                    BatchTaskResult(
                        index=submission.index,
                        task_id=submission.task_id,
                        submitted=True,
                        artifacts=artifacts,
                        results=task_results,
                    )
                )
            except Exception as exc:
                results.append(
                    BatchTaskResult(
                        index=submission.index,
                        task_id=submission.task_id,
                        submitted=True,
                        error=exc,
                    )
                )
            finally:
                await self._forget_result_if_needed(handle)

        return results

    async def batch_sync_detailed(
        self,
        tasks: LegacyBatchTasks,
        *,
        priority: Optional[str] = None,
        timeout: Optional[int] = None,
        cancel_key: Optional[str] = None,
        persistent_output: bool = False,
    ) -> List[BatchTaskResult]:
        submissions = await self.batch_submit(
            tasks,
            priority=priority,
            timeout=timeout,
            cancel_key=cancel_key,
            persistent_output=persistent_output,
        )
        return await self.batch_collect(submissions, timeout=timeout)

    async def batch_sync(
        self,
        tasks: LegacyBatchTasks,
        *,
        priority: Optional[str] = None,
        timeout: Optional[int] = None,
        cancel_key: Optional[str] = None,
        persistent_output: bool = False,
    ) -> List[LegacyTaskResult]:
        detailed_results = await self.batch_sync_detailed(
            tasks,
            priority=priority,
            timeout=timeout,
            cancel_key=cancel_key,
            persistent_output=persistent_output,
        )
        return [result.to_legacy_tuple() for result in detailed_results]

    def get_queue_info(self) -> Dict[str, str]:
        return self.queue_names.copy()

    async def probe(self, timeout: int = 10) -> HealthCheckResult:
        try:
            handle = await self._send_task_async(
                input_artifacts=[],
                options={"test": True},
                priority="low",
            )
        except Exception as exc:
            return HealthCheckResult(
                submit_ok=False,
                retrieval_ok=False,
                response_ok=False,
                error=str(exc),
            )

        try:
            raw_result = await self.transport.get_result(handle.async_result, timeout=timeout)
            response_ok, probe_error = inspect_probe_result(raw_result)
            return HealthCheckResult(
                submit_ok=True,
                retrieval_ok=True,
                response_ok=response_ok,
                task_id=handle.task_id,
                error=probe_error,
            )
        except Exception as exc:
            return HealthCheckResult(
                submit_ok=True,
                retrieval_ok=False,
                response_ok=False,
                task_id=handle.task_id,
                error=str(exc),
            )
        finally:
            await self._forget_result_if_needed(handle)

    async def ping(self, timeout: int = 10) -> bool:
        return (await self.probe(timeout=timeout)).submit_ok



def create_client(
    operator_name: str, operator_version: str, **kwargs: Any
) -> OperatorClient:
    """Convenience constructor for ``OperatorClient``."""

    return OperatorClient(operator_name, operator_version, **kwargs)


def create_async_client(
    operator_name: str, operator_version: str, **kwargs: Any
) -> AsyncOperatorClient:
    """Convenience constructor for ``AsyncOperatorClient``."""

    return AsyncOperatorClient(operator_name, operator_version, **kwargs)


__all__ = [
    "CeleryConfig",
    "ClientConfig",
    "LegacyBatchTasks",
    "LegacyTaskResult",
    "OperatorClient",
    "HealthCheckResult",
    "SubmittedTask",
    "BatchTaskResult",
    "AsyncOperatorClient",
    "AsyncTaskHandle",
    "AsyncSubmittedTask",
    "create_client",
    "create_async_client",
    "get_celery_config",
    "get_default_celery_config",
    "resolve_client_config",
]
