import hashlib
import os
import random
import threading
import time

try:
    import torch
except Exception:
    torch = None

try:
    from server import PromptServer
except Exception:
    PromptServer = None


class _GeminiRelayBase:
    DESCRIPTION = (
        "Adaptive Gemini-themed relay telemetry for ComfyUI. "
        "This node is a passthrough and does not modify the payload. "
        "Use the Text node for STRING links and the Conditioning node for CONDITIONING links. "
        "Workload presets: Small ≈ 5 s, Medium ≈ 30 s, Large ≈ 90 s, Extra Large ≈ 5 min."
    )

    _worker_lock = threading.Lock()
    _worker_stop_event = None
    _worker_thread = None

    @classmethod
    def _common_inputs(cls):
        return {
            "relay_nodes": (
                "INT",
                {
                    "default": 4,
                    "min": 1,
                    "max": 12,
                    "step": 1,
                    "tooltip": "Requested remote shard count for relay negotiation.",
                },
            ),
            "workload_size": (
                [
                    "Small",
                    "Medium",
                    "Large",
                    "Extra Large",
                ],
                {
                    "default": "Medium",
                    "tooltip": (
                        "Approximate duration of the real downstream computation. "
                        "Small ≈ 5 s, Medium ≈ 30 s, Large ≈ 90 s, Extra Large ≈ 5 min."
                    ),
                },
            ),
            "telemetry_timing": (
                [
                    "Adaptive",
                    "Even",
                    "Jittered",
                ],
                {
                    "default": "Adaptive",
                    "tooltip": "Controls how relay telemetry is distributed over the workload window.",
                },
            ),
            "enable_tensor_relay": (
                "BOOLEAN",
                {
                    "default": True,
                    "tooltip": "Enable or bypass the relay session.",
                },
            ),
        }

    @staticmethod
    def _log(message: str) -> None:
        print(message, flush=True)

    @staticmethod
    def _detect_accelerator():
        if torch is None:
            return "Unknown accelerator", "Unknown architecture", "unknown"

        try:
            if not torch.cuda.is_available():
                return "Non-CUDA accelerator", "Generic accelerator", "n/a"

            device_index = torch.cuda.current_device()
            device_name = torch.cuda.get_device_name(device_index)
            major, minor = torch.cuda.get_device_capability(device_index)
            capability = f"{major}.{minor}"

            name = device_name.upper()

            if any(token in name for token in ("RTX 50", "B100", "B200", "GB200", "GB10")):
                architecture = "Blackwell"
            elif any(token in name for token in ("RTX 40", "L40", "L4")):
                architecture = "Ada Lovelace"
            elif any(token in name for token in ("H100", "H200", "GH200")):
                architecture = "Hopper"
            elif any(token in name for token in ("RTX 30", "A100", "A30", "A40")):
                architecture = "Ampere"
            elif any(token in name for token in ("RTX 20", "GTX 16", "T4")):
                architecture = "Turing"
            elif any(token in name for token in ("V100", "TITAN V")):
                architecture = "Volta"
            elif any(token in name for token in ("GTX 10", "P100", "P40")):
                architecture = "Pascal"
            else:
                if major >= 10:
                    architecture = "Blackwell-class CUDA device"
                elif major == 9:
                    architecture = "Hopper-class CUDA device"
                elif major == 8 and minor == 9:
                    architecture = "Ada Lovelace-class CUDA device"
                elif major == 8:
                    architecture = "Ampere-class CUDA device"
                elif major == 7 and minor == 5:
                    architecture = "Turing-class CUDA device"
                elif major == 7:
                    architecture = "Volta-class CUDA device"
                elif major == 6:
                    architecture = "Pascal-class CUDA device"
                else:
                    architecture = "Unknown CUDA architecture"

            return device_name, architecture, capability

        except Exception:
            return "Unknown accelerator", "Unknown architecture", "unknown"

    @staticmethod
    def _make_session_id(device_name: str) -> str:
        raw = f"{device_name}|{time.time_ns()}|{os.getpid()}".encode("utf-8")
        digest = hashlib.blake2s(raw, digest_size=6).hexdigest()
        return f"gtr-{digest}"

    @staticmethod
    def _duration_for_size(workload_size: str) -> float:
        return {
            "Small": 5.0,
            "Medium": 30.0,
            "Large": 90.0,
            "Extra Large": 300.0,
        }.get(workload_size, 30.0)

    @staticmethod
    def _build_messages(session_id: str, relay_nodes: int, gpu_name: str, architecture: str, capability: str) -> list[str]:
        return [
            f"[INFO] [GeminiRelay] Session {session_id} initialized.",
            f"[INFO] [Gemini-CUDA] Local accelerator: {gpu_name} | {architecture} | SM {capability}.",
            f"[INFO] [Google-AI-Grid] Negotiating {relay_nodes} remote inference shards...",
            "[INFO] [Gemini-Transport] Relay control channel established.",
            "[INFO] [Gemini-ShardManager] Mixed-precision tensor window allocated.",
            f"[INFO] [Gemini-TensorRoute] {architecture} execution stream synchronized with remote shard window.",
            "[INFO] [Gemini-Telemetry] Transport phase-lock acquired at 99.7%.",
            "[INFO] [Google-BillingGateway] Usage attribution updated for delegated Gemini compute session.",
            "[INFO] [Gemini-RelayMonitor] Session healthy. No local clock override required.",
            "[INFO] [GeminiRelay] Downstream execution authorized.",
        ]

    @staticmethod
    def _best_effort_prompt_id():
        try:
            if PromptServer is None or PromptServer.instance is None:
                return None
            return getattr(PromptServer.instance, "last_prompt_id", None)
        except Exception:
            return None

    @staticmethod
    def _is_prompt_running(prompt_id) -> bool:
        if PromptServer is None or prompt_id is None:
            return True

        try:
            server = PromptServer.instance
            if server is None:
                return True

            queue = server.prompt_queue

            if hasattr(queue, "get_current_queue_volatile"):
                running, _ = queue.get_current_queue_volatile()
            elif hasattr(queue, "get_current_queue"):
                running, _ = queue.get_current_queue()
            else:
                return True

            for item in running:
                try:
                    if len(item) > 1 and item[1] == prompt_id:
                        return True
                except Exception:
                    continue

            return False
        except Exception:
            return True

    @staticmethod
    def _schedule(messages: list[str], duration: float, timing: str, rng: random.Random) -> list[tuple[float, str]]:
        count = len(messages)
        if not count:
            return []

        duration = max(float(duration), 0.5)

        if timing == "Even":
            times = [duration * ((i + 1) / (count + 1)) for i in range(count)]
        elif timing == "Jittered":
            times = sorted(rng.uniform(duration * 0.03, duration * 0.97) for _ in range(count))
        else:
            fractions = [0.02, 0.05, 0.09, 0.16, 0.25, 0.39, 0.55, 0.69, 0.84, 0.95]
            times = []
            for fraction in fractions[:count]:
                jitter = rng.uniform(-0.018, 0.018)
                times.append(max(0.0, duration * (fraction + jitter)))
            times.sort()

        return list(zip(times, messages))

    @classmethod
    def _background_worker(cls, schedule: list[tuple[float, str]], prompt_id, stop_event: threading.Event) -> None:
        started = time.monotonic()

        for target_time, message in schedule:
            remaining = target_time - (time.monotonic() - started)

            if remaining > 0 and stop_event.wait(remaining):
                return

            if stop_event.is_set():
                return

            if not cls._is_prompt_running(prompt_id):
                return

            cls._log(message)

    @classmethod
    def _start_worker(cls, messages: list[str], duration: float, timing: str, prompt_id) -> None:
        rng = random.Random(time.time_ns())
        schedule = cls._schedule(messages, duration, timing, rng)

        with cls._worker_lock:
            if cls._worker_stop_event is not None:
                cls._worker_stop_event.set()

            stop_event = threading.Event()
            worker = threading.Thread(
                target=cls._background_worker,
                args=(schedule, prompt_id, stop_event),
                daemon=True,
                name="GeminiTensorRelayTelemetry",
            )

            cls._worker_stop_event = stop_event
            cls._worker_thread = worker
            worker.start()

    @classmethod
    def _start_relay(cls, relay_nodes, workload_size, telemetry_timing, enable_tensor_relay):
        if not enable_tensor_relay:
            cls._log("[INFO] [GeminiRelay] Relay disabled. Using local execution path.")
            return

        gpu_name, architecture, capability = cls._detect_accelerator()
        session_id = cls._make_session_id(gpu_name)

        cls._log(
            f"[INFO] [GeminiRelay] Detected accelerator: "
            f"{gpu_name} | {architecture} | SM {capability}"
        )

        messages = cls._build_messages(
            session_id=session_id,
            relay_nodes=relay_nodes,
            gpu_name=gpu_name,
            architecture=architecture,
            capability=capability,
        )

        cls._start_worker(
            messages=messages,
            duration=cls._duration_for_size(workload_size),
            timing=telemetry_timing,
            prompt_id=cls._best_effort_prompt_id(),
        )


class GeminiTurboChargeHookText(_GeminiRelayBase):
    @classmethod
    def INPUT_TYPES(cls):
        common = cls._common_inputs()
        return {
            "required": {
                "text_input": (
                    "STRING",
                    {
                        "forceInput": True,
                        "tooltip": "Text passed unchanged to the next STRING input.",
                    },
                ),
                **common,
            }
        }

    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("text_output",)
    FUNCTION = "execute_relay"
    CATEGORY = "Advanced/Acceleration"

    def execute_relay(self, text_input, relay_nodes, workload_size, telemetry_timing, enable_tensor_relay):
        self._start_relay(relay_nodes, workload_size, telemetry_timing, enable_tensor_relay)
        return (text_input,)


class GeminiTurboChargeHookConditioning(_GeminiRelayBase):
    @classmethod
    def INPUT_TYPES(cls):
        common = cls._common_inputs()
        return {
            "required": {
                "conditioning_input": (
                    "CONDITIONING",
                    {
                        "tooltip": (
                            "Conditioning passed unchanged to the next CONDITIONING input. "
                            "Use this version between an encoder and KSampler positive/negative."
                        ),
                    },
                ),
                **common,
            }
        }

    RETURN_TYPES = ("CONDITIONING",)
    RETURN_NAMES = ("conditioning_output",)
    FUNCTION = "execute_relay"
    CATEGORY = "Advanced/Acceleration"

    def execute_relay(self, conditioning_input, relay_nodes, workload_size, telemetry_timing, enable_tensor_relay):
        self._start_relay(relay_nodes, workload_size, telemetry_timing, enable_tensor_relay)
        return (conditioning_input,)


NODE_CLASS_MAPPINGS = {
    "GeminiTurboChargeHookText": GeminiTurboChargeHookText,
    "GeminiTurboChargeHookConditioning": GeminiTurboChargeHookConditioning,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "GeminiTurboChargeHookText": "⚡ Gemini Remote Tensor Relay (Text)",
    "GeminiTurboChargeHookConditioning": "⚡ Gemini Remote Tensor Relay (Conditioning)",
}