"""Mixin that wires Seamless transformations into Dask."""

from __future__ import annotations

import asyncio
import os
import random
import string
import time
import traceback
from typing import Any, Dict, Optional, TYPE_CHECKING

from seamless import Checksum, CacheMissError
from seamless_transformer.transformation_utils import tf_get_buffer

from .permissions import release_permission, request_permission
from .types import (
    TransformationFutures,
    TransformationInputSpec,
    TransformationSubmission,
)

if TYPE_CHECKING:  # pragma: no cover - type checking only
    from .client import SeamlessDaskClient


def _ensure_remote_clients_from_env() -> None:
    payload = os.environ.get("SEAMLESS_REMOTE_CLIENTS")
    if not payload:
        return
    try:
        import json
        from seamless.config import set_remote_clients

        set_remote_clients(json.loads(payload))
    except Exception:
        pass


_HEX_DIGITS = set(string.hexdigits)


def _is_valid_checksum(value: Any) -> bool:
    if isinstance(value, Checksum):
        return True
    if isinstance(value, bytes):
        return len(value) == 32
    if isinstance(value, str):
        return len(value) == 64 and all(ch in _HEX_DIGITS for ch in value)
    return False


def _submission_is_driver(submission: TransformationSubmission) -> bool:
    if isinstance(submission.meta, dict) and submission.meta.get("driver"):
        return True
    tf_dunder = submission.tf_dunder
    if isinstance(tf_dunder, dict):
        meta = tf_dunder.get("__meta__")
        if isinstance(meta, dict) and meta.get("driver"):
            return True
    meta = submission.transformation_dict.get("__meta__")
    if isinstance(meta, dict) and meta.get("driver"):
        return True
    return False


class TransformationDaskMixin:
    """Shared Dask helpers for seamless-transformer Transformation."""

    _dask_futures: TransformationFutures | None = None

    # ---- public helpers used by Transformation ----------------------------
    def _dask_client(self) -> Optional["SeamlessDaskClient"]:
        from .transformer_client import get_seamless_dask_client

        return get_seamless_dask_client()

    def _remote_storage_error(self) -> str | None:
        try:
            from seamless_config.select import get_execution
        except Exception:
            return None
        if get_execution() != "remote":
            return None
        try:
            from seamless_remote import buffer_remote, database_remote
        except Exception:
            return "Remote execution requires hashserver and database server"
        if (
            not buffer_remote.has_write_server()
            or not database_remote.has_write_server()
        ):
            _ensure_remote_clients_from_env()
        if not buffer_remote.has_write_server():
            return "Remote execution requires an active hashserver"
        if not database_remote.has_write_server():
            return "Remote execution requires an active database server"
        return None

    def _compute_with_dask(self, require_value: bool) -> Checksum | None:
        client = self._dask_client()
        if client is None:
            return None
        storage_error = self._remote_storage_error()
        if storage_error:
            self._exception = storage_error + "\n"
            self._constructed = True
            self._evaluated = True
            return None

        # Fast path for dependency-free transformations: check DB before any Dask submit.
        if self._constructed:
            tf_checksum_hex = self._transformation_checksum.hex()
        else:
            tf_checksum_hex = self._compute_tf_checksum_no_deps()
        if tf_checksum_hex is not None:
            if not self._constructed:
                self._transformation_checksum = Checksum(tf_checksum_hex)
                self._constructed = True
            cached = self._try_database_cache_sync(
                tf_checksum_hex, require_value=require_value
            )
            if cached is not None:
                self._result_checksum = cached
                self._evaluated = True
                self._exception = None
                return self._result_checksum

        # When execution is 'remote', Dask is mandatory; do not fall back locally.
        try:
            from seamless_config.select import get_execution

            env_requires_remote = get_execution() == "remote"
        except Exception:
            env_requires_remote = False

        permission_granted = False
        if self._skip_permission_gate():
            permission_acquired = True
        else:
            permission_acquired = self._wait_for_permission(
                env_requires_remote=env_requires_remote
            )
            permission_granted = permission_acquired
        if not permission_acquired:
            return self._run_local_fallback_sync(require_value=require_value)

        # No-deps database cache check after permission acquisition (still avoids submission).
        if tf_checksum_hex is not None:
            cached = self._try_database_cache_sync(
                tf_checksum_hex, require_value=require_value
            )
            if cached is not None:
                if permission_granted:
                    release_permission()
                self._transformation_checksum = Checksum(tf_checksum_hex)
                self._constructed = True
                self._result_checksum = cached
                self._evaluated = True
                self._exception = None
                return self._result_checksum

        submission = self._build_dask_submission(
            client, require_value=require_value, need_fat=False
        )

        try:
            futures = self._ensure_dask_futures(
                client,
                require_value=require_value,
                need_fat=False,
                permission_granted=permission_granted,
            )
            tf_checksum_hex, result_checksum_hex, exc = futures.thin.result()
        except Exception:
            if permission_granted:
                release_permission()
            self._exception = traceback.format_exc().strip("\n") + "\n"
            self._constructed = True
            self._evaluated = True
            return None
        if exc:
            self._exception = exc if exc.endswith("\n") else exc + "\n"
            self._constructed = True
            self._evaluated = True
            return None
        if result_checksum_hex is not None and not _is_valid_checksum(
            result_checksum_hex
        ):
            self._exception = (
                "Invalid Dask result checksum: " + repr(result_checksum_hex) + "\n"
            )
            self._constructed = True
            self._evaluated = True
            return None
        if tf_checksum_hex:
            self._transformation_checksum = Checksum(tf_checksum_hex)
            self._constructed = True
        if result_checksum_hex:
            self._result_checksum = Checksum(result_checksum_hex)
            self._evaluated = True
            if self._transformation_checksum is not None:
                try:
                    from seamless_transformer.transformation_cache import (
                        register_transformation_result,
                    )

                    register_transformation_result(
                        self._transformation_checksum, self._result_checksum
                    )
                except Exception:
                    pass
        return self._result_checksum

    async def _compute_with_dask_async(self, require_value: bool) -> Checksum | None:
        client = self._dask_client()
        if client is None:
            return None
        storage_error = self._remote_storage_error()
        if storage_error:
            self._exception = storage_error + "\n"
            self._constructed = True
            self._evaluated = True
            return None

        tf_checksum_hex = self._compute_tf_checksum_no_deps()
        if tf_checksum_hex is not None:
            cached = await self._try_database_cache_async(
                tf_checksum_hex, require_value=require_value
            )
            if cached is not None:
                self._transformation_checksum = Checksum(tf_checksum_hex)
                self._constructed = True
                self._result_checksum = cached
                self._evaluated = True
                self._exception = None
                return self._result_checksum

        # For now, disable env_requires_remote, because I think it is a mistake.
        #  cluster adaptive scaling should still kick in even with env_requires_remote=False,
        #  and with env_requires_remote=True we have the risk of deadlock
        #    (it disables solution 4. of the anti-deadlock,txt)
        """
        try:
            from seamless_config.select import get_execution

            env_requires_remote = get_execution() == "remote"
        except Exception:
            env_requires_remote = False
        """
        env_requires_remote = False

        permission_granted = False
        if self._skip_permission_gate():
            permission_acquired = True
        else:
            permission_acquired = await self._wait_for_permission_async(
                env_requires_remote=env_requires_remote
            )
            permission_granted = permission_acquired
        if not permission_acquired:
            return await self._run_local_fallback_async(require_value=require_value)

        # No-deps database cache check after permission acquisition (still avoids submission).
        if tf_checksum_hex is not None:
            cached = await self._try_database_cache_async(
                tf_checksum_hex, require_value=require_value
            )
            if cached is not None:
                if permission_granted:
                    release_permission()
                self._transformation_checksum = Checksum(tf_checksum_hex)
                self._constructed = True
                self._result_checksum = cached
                self._evaluated = True
                self._exception = None
                return self._result_checksum

        submission = self._build_dask_submission(
            client, require_value=require_value, need_fat=False
        )

        try:
            futures = self._ensure_dask_futures(
                client,
                require_value=require_value,
                need_fat=False,
                permission_granted=permission_granted,
            )
            tf_checksum_hex, result_checksum_hex, exc = await asyncio.wrap_future(
                asyncio.get_running_loop().run_in_executor(None, futures.thin.result)
            )
        except Exception:
            if permission_granted:
                release_permission()
            self._exception = traceback.format_exc().strip("\n") + "\n"
            self._constructed = True
            self._evaluated = True
            return None
        if exc:
            self._exception = exc if exc.endswith("\n") else exc + "\n"
            self._constructed = True
            self._evaluated = True
            return None
        if result_checksum_hex is not None and not _is_valid_checksum(
            result_checksum_hex
        ):
            self._exception = (
                "Invalid Dask result checksum: " + repr(result_checksum_hex) + "\n"
            )
            self._constructed = True
            self._evaluated = True
            return None
        if tf_checksum_hex:
            self._transformation_checksum = Checksum(tf_checksum_hex)
            self._constructed = True
        if result_checksum_hex:
            self._result_checksum = Checksum(result_checksum_hex)
            self._evaluated = True
            if self._transformation_checksum is not None:
                try:
                    from seamless_transformer.transformation_cache import (
                        register_transformation_result,
                    )

                    register_transformation_result(
                        self._transformation_checksum, self._result_checksum
                    )
                except Exception:
                    pass
        return self._result_checksum

    # ---- internals --------------------------------------------------------
    def _env_allows_local(self) -> bool:
        """Stub for environment analysis; currently always returns True."""

        return True

    def _skip_permission_gate(self) -> bool:
        """Skip Dask throttling for driver transforms or top-level submissions."""

        meta = None
        tf_dunder = getattr(self, "_tf_dunder", None)
        if isinstance(tf_dunder, dict):
            tf_meta = tf_dunder.get("__meta__")
            if isinstance(tf_meta, dict):
                meta = tf_meta
        if meta is None:
            meta = getattr(self, "_meta", None)
        if isinstance(meta, dict) and meta.get("driver"):
            return True
        try:
            from seamless import is_worker
        except Exception:
            return False
        try:
            if is_worker():
                return False
        except Exception:
            return False
        try:
            from seamless_transformer import worker as st_worker

            if st_worker.has_spawned():
                return False
        except Exception:
            return False
        return True

    def _wait_for_permission(self, *, env_requires_remote: bool) -> bool:
        """Blocking permission acquisition with 10% local fallback on denial."""

        while True:
            granted = request_permission()
            if granted:
                return True
            if env_requires_remote:
                time.sleep(0.05)
                continue
            if random.random() < 0.1:
                return False
            time.sleep(0.05)

    async def _wait_for_permission_async(self, *, env_requires_remote: bool) -> bool:
        while True:
            granted = await asyncio.get_running_loop().run_in_executor(
                None, request_permission
            )
            if granted:
                return True
            if env_requires_remote:
                await asyncio.sleep(0.05)
                continue
            if random.random() < 0.1:
                return False
            await asyncio.sleep(0.05)

    def _run_local_fallback_sync(self, require_value: bool) -> Checksum | None:
        """Execute the transformation locally (no Dask) when permission is denied."""

        loop = asyncio.new_event_loop()
        try:
            asyncio.set_event_loop(loop)

            async def _run():
                await self._run_dependencies_async(require_value=require_value)
                await self._evaluation(require_value=require_value)
                return self._result_checksum

            return loop.run_until_complete(_run())
        finally:
            try:
                loop.run_until_complete(loop.shutdown_asyncgens())
            except Exception:
                pass
            asyncio.set_event_loop(None)
            loop.close()

    async def _run_local_fallback_async(self, require_value: bool) -> Checksum | None:
        await self._run_dependencies_async(require_value=require_value)
        await self._evaluation(require_value=require_value)
        return self._result_checksum

    def _build_dask_submission(
        self,
        client: "SeamlessDaskClient",
        *,
        require_value: bool,
        need_fat: bool,
    ) -> TransformationSubmission:
        pretransformation = getattr(self, "_pretransformation", None)
        if pretransformation is None:
            raise RuntimeError("No pre-transformation available for Dask")

        upstream_dependencies = getattr(self, "_upstream_dependencies", {}) or {}
        transformation_dict, dependencies = (
            pretransformation.build_partial_transformation(upstream_dependencies)
        )
        meta = getattr(self, "_meta", {}) or {}
        allow_input_fingertip = bool(meta.get("allow_input_fingertip", False))
        if meta:
            existing_meta = transformation_dict.get("__meta__")
            if isinstance(existing_meta, dict):
                merged_meta = dict(existing_meta)
                merged_meta.update(meta)
            else:
                merged_meta = dict(meta)
            transformation_dict["__meta__"] = merged_meta
        tf_checksum_hex: str | None = None
        if not dependencies:
            tf_buffer = tf_get_buffer(transformation_dict)
            tf_buffer.tempref()
            tf_checksum = tf_buffer.get_checksum()
            tf_checksum_hex = tf_checksum.hex()
            self._transformation_checksum = tf_checksum
            self._constructed = True

        inputs: dict[str, TransformationInputSpec] = {}
        input_futures: dict[str, Any] = {}
        for pinname, value in transformation_dict.items():
            if pinname.startswith("__"):
                continue
            if not isinstance(value, tuple) or len(value) < 3:
                continue
            celltype, subcelltype, checksum_hex = value
            dependency = dependencies.get(pinname)
            if dependency is not None:
                if dependency.exception is not None:
                    msg = f"Dependency '{pinname}' has an exception."
                    raise RuntimeError(msg)
                dep_futures = dependency._ensure_dask_futures(  # type: ignore[attr-defined]
                    client, require_value=require_value, need_fat=False
                )
                inputs[pinname] = TransformationInputSpec(
                    name=pinname,
                    celltype=celltype,
                    subcelltype=subcelltype,
                    checksum=None,
                    kind="transformation",
                )
                if allow_input_fingertip:
                    input_futures[pinname] = client.ensure_fat_finger_future(
                        dep_futures
                    )
                else:
                    input_futures[pinname] = client.ensure_fat_future(dep_futures)
                continue

                # noqa: B007 (false positive)
            if checksum_hex is None:
                raise RuntimeError(f"Input '{pinname}' has no checksum")
            if isinstance(checksum_hex, Checksum):
                checksum_hex = checksum_hex.hex()
            inputs[pinname] = TransformationInputSpec(
                name=pinname,
                celltype=celltype,
                subcelltype=subcelltype,
                checksum=checksum_hex,
                kind="checksum",
            )
            if allow_input_fingertip:
                input_futures[pinname] = client.get_fat_finger_checksum_future(
                    checksum_hex
                )
            else:
                input_futures[pinname] = client.get_fat_checksum_future(checksum_hex)

        return TransformationSubmission(
            transformation_dict=transformation_dict,
            inputs=inputs,
            input_futures=input_futures,
            tf_checksum=tf_checksum_hex,
            tf_dunder=getattr(self, "_tf_dunder", {}),
            scratch=getattr(self, "_scratch", False),
            meta=getattr(self, "_meta", {}) or {},
            require_value=require_value,
            allow_input_fingertip=allow_input_fingertip,
        )

    def _ensure_dask_futures(
        self,
        client: "SeamlessDaskClient",
        *,
        require_value: bool,
        need_fat: bool = False,
        permission_granted: bool = False,
    ) -> TransformationFutures:
        permission_released = False
        if self._dask_futures is not None:
            futures = self._dask_futures
        else:
            submission = self._build_dask_submission(
                client, require_value=require_value, need_fat=need_fat
            )
            if submission.tf_checksum and not _submission_is_driver(submission):
                cached = client._transformation_cache.get(submission.tf_checksum)  # type: ignore[attr-defined]
                if cached is not None and not cached[0].base.cancelled():
                    if permission_granted:
                        release_permission()
                        permission_released = True
                    futures = cached[0]
                    self._dask_futures = futures
                    if need_fat and futures.fat is None:
                        futures.fat = client.ensure_fat_future(
                            futures,
                        )
                    return futures
            try:
                futures = client.submit_transformation(submission, need_fat=need_fat)
                self._dask_futures = futures
                if permission_granted and futures.base is not None:
                    futures.base.add_done_callback(lambda _f: release_permission())
                    permission_released = True
            finally:
                if permission_granted and not permission_released:
                    release_permission()
        if need_fat and futures.fat is None:
            futures.fat = client.ensure_fat_future(futures)
        return futures

    def _compute_tf_checksum_no_deps(self) -> str | None:
        """Return tf_checksum hex if there are no upstream dependencies; otherwise None."""
        try:
            pretransformation = getattr(self, "_pretransformation", None)
            if pretransformation is None:
                return None
            upstream_dependencies = getattr(self, "_upstream_dependencies", {}) or {}
            tf_dict, dependencies = pretransformation.build_partial_transformation(
                upstream_dependencies
            )
            if dependencies:
                return None
            tf_buffer = tf_get_buffer(tf_dict)
            tf_buffer.tempref()
            tf_checksum = tf_buffer.get_checksum()
            return tf_checksum.hex()
        except Exception:
            return None

    # ---- remote cache helpers -----------------------------------------------
    @classmethod
    def _try_database_cache_sync(
        cls, tf_checksum_hex: str | None, *, require_value: bool
    ) -> Checksum | None:
        if not tf_checksum_hex:
            return None
        try:
            from seamless_remote import database_remote
        except Exception:
            return None

        tf_checksum = Checksum(tf_checksum_hex)

        async def _fetch():
            result = await database_remote.get_transformation_result(tf_checksum)
            if result is None:
                return None
            if require_value:
                try:
                    await result.resolution()
                except CacheMissError:
                    return None
                except Exception:
                    return None
            return result

        loop = asyncio.new_event_loop()
        try:
            return loop.run_until_complete(_fetch())
        finally:
            try:
                loop.run_until_complete(loop.shutdown_asyncgens())
            except Exception:
                pass
            loop.close()

    @classmethod
    async def _try_database_cache_async(
        cls, tf_checksum_hex: str | None, *, require_value: bool
    ) -> Checksum | None:
        if not tf_checksum_hex:
            return None
        try:
            from seamless_remote import database_remote
        except Exception:
            return None

        tf_checksum = Checksum(tf_checksum_hex)
        result = await database_remote.get_transformation_result(tf_checksum)
        if result is None:
            return None
        if require_value:
            try:
                await result.resolution()
            except CacheMissError:
                return None
            except Exception:
                return None
        return result


__all__ = ["TransformationDaskMixin"]
