import asyncio
import json
import time
import uuid
from collections import deque

import aiohttp

import execution
import server

from .utils.config import load_config
from .utils.logging import debug_log, log


prompt_server = server.PromptServer.instance


def ensure_distributed_state():
    """Ensure prompt_server has the state used by distributed queue orchestration."""
    if not hasattr(prompt_server, "distributed_pending_jobs"):
        prompt_server.distributed_pending_jobs = {}
    if not hasattr(prompt_server, "distributed_jobs_lock"):
        prompt_server.distributed_jobs_lock = asyncio.Lock()


async def _get_client_session():
    """Get or create aiohttp client session (shared with distributed.py)."""
    if not hasattr(prompt_server, "_distributed_session"):
        prompt_server._distributed_session = aiohttp.ClientSession()
    return prompt_server._distributed_session


def _deepcopy_prompt(prompt_obj):
    """Return a deep copy of the workflow/prompt dictionary."""
    return json.loads(json.dumps(prompt_obj))


def _iter_prompt_nodes(prompt_obj):
    for node_id, node in prompt_obj.items():
        if isinstance(node, dict):
            yield str(node_id), node


def _find_nodes_by_class(prompt_obj, class_name):
    nodes = []
    for node_id, node in _iter_prompt_nodes(prompt_obj):
        if node.get("class_type") == class_name:
            nodes.append(node_id)
    return nodes


def _has_upstream_node(prompt_obj, start_node_id, target_class):
    """Depth-first search to find if start node depends on target_class."""
    visited = set()
    stack = [start_node_id]

    while stack:
        node_id = stack.pop()
        if node_id in visited:
            continue
        visited.add(node_id)
        node = prompt_obj.get(node_id)
        if not node:
            continue
        inputs = node.get("inputs", {})
        for value in inputs.values():
            if isinstance(value, list) and len(value) == 2:
                upstream_id = str(value[0])
                upstream_node = prompt_obj.get(upstream_id)
                if not upstream_node:
                    continue
                if upstream_node.get("class_type") == target_class:
                    return True
                stack.append(upstream_id)
    return False


def _find_downstream_nodes(prompt_obj, start_ids):
    """Return all nodes reachable downstream from the provided IDs."""
    adjacency = {}
    for node_id, node in _iter_prompt_nodes(prompt_obj):
        inputs = node.get("inputs", {})
        for value in inputs.values():
            if isinstance(value, list) and len(value) == 2:
                source_id = str(value[0])
                adjacency.setdefault(source_id, set()).add(str(node_id))

    connected = set(start_ids)
    queue = deque(start_ids)
    while queue:
        current = queue.popleft()
        for dependent in adjacency.get(current, ()):  # pragma: no branch - simple iteration
            if dependent not in connected:
                connected.add(dependent)
                queue.append(dependent)
    return connected


def _create_numeric_id_generator(prompt_obj):
    """Return a closure that yields new numeric string IDs."""
    max_id = 0
    for node_id in prompt_obj.keys():
        try:
            numeric = int(node_id)
        except (TypeError, ValueError):
            continue
        max_id = max(max_id, numeric)

    counter = max_id

    def _next_id():
        nonlocal counter
        counter += 1
        return str(counter)

    return _next_id


def _prepare_delegate_master_prompt(prompt_obj, collector_ids):
    """Prune master prompt so it only executes post-collector nodes in delegate mode."""
    downstream = _find_downstream_nodes(prompt_obj, collector_ids)
    nodes_to_keep = set(collector_ids)
    nodes_to_keep.update(downstream)

    pruned_prompt = {}
    for node_id in nodes_to_keep:
        node = prompt_obj.get(node_id)
        if node is not None:
            pruned_prompt[node_id] = json.loads(json.dumps(node))

    pruned_ids = set(pruned_prompt.keys())
    for node_id, node in pruned_prompt.items():
        inputs = node.get("inputs")
        if not inputs:
            continue
        for input_name, input_value in list(inputs.items()):
            if isinstance(input_value, list) and len(input_value) == 2:
                source_id = str(input_value[0])
                if source_id not in pruned_ids:
                    inputs.pop(input_name, None)
                    debug_log(
                        f"Removed upstream reference '{input_name}' from node {node_id} for delegate-only master prompt."
                    )

    next_id = _create_numeric_id_generator(pruned_prompt)
    for collector_id in collector_ids:
        collector_entry = pruned_prompt.get(collector_id)
        if not collector_entry:
            continue
        placeholder_id = next_id()
        pruned_prompt[placeholder_id] = {
            "class_type": "DistributedEmptyImage",
            "inputs": {
                "height": 64,
                "width": 64,
                "channels": 3,
            },
            "_meta": {
                "title": "Distributed Empty Image (auto-added)",
            },
        }
        collector_entry.setdefault("inputs", {})["images"] = [placeholder_id, 0]
        debug_log(
            f"Inserted placeholder node {placeholder_id} for collector {collector_id} in delegate-only master prompt."
        )

    return pruned_prompt


async def _ensure_distributed_queue(job_id):
    """Ensure a queue exists for the given distributed job ID."""
    ensure_distributed_state()
    async with prompt_server.distributed_jobs_lock:
        if job_id not in prompt_server.distributed_pending_jobs:
            prompt_server.distributed_pending_jobs[job_id] = asyncio.Queue()


def _generate_job_id_map(prompt_obj, prefix):
    """Create stable per-node job IDs for distributed nodes."""
    job_map = {}
    distributed_nodes = _find_nodes_by_class(prompt_obj, "DistributedCollector") + _find_nodes_by_class(
        prompt_obj, "UltimateSDUpscaleDistributed"
    )
    for node_id in distributed_nodes:
        job_map[node_id] = f"{prefix}_{node_id}"
    return job_map


def _resolve_enabled_workers(config, requested_ids=None):
    """Return a list of worker configs that should participate."""
    workers = []
    for worker in config.get("workers", []):
        worker_id = str(worker.get("id") or "").strip()
        if not worker_id:
            continue

        if requested_ids is not None:
            if worker_id not in requested_ids:
                continue
        elif not worker.get("enabled", False):
            continue

        workers.append(
            {
                "id": worker_id,
                "name": worker.get("name", worker_id),
                "host": worker.get("host"),
                "port": int(worker.get("port", worker.get("listen_port", 8188)) or 8188),
                "type": worker.get("type", "local"),
            }
        )
    return workers


def _resolve_master_url():
    """Best-effort reconstruction of the master's public URL."""
    cfg = load_config()
    master_cfg = cfg.get("master", {}) or {}
    configured_host = (master_cfg.get("host") or "").strip()
    configured_port = master_cfg.get("port")
    default_port = getattr(prompt_server, "port", 8188) or 8188
    port = int(configured_port or default_port)

    def _needs_https(hostname):
        hostname = hostname.lower()
        https_domains = (
            ".proxy.runpod.net",
            ".ngrok-free.app",
            ".ngrok-free.dev",
            ".ngrok.io",
        )
        return any(hostname.endswith(suffix) for suffix in https_domains)

    if configured_host:
        if configured_host.startswith(("http://", "https://")):
            return configured_host.rstrip("/")

        host = configured_host
        scheme = "https" if _needs_https(host) or port == 443 else "http"
        default_port_for_scheme = 443 if scheme == "https" else 80
        # For ngrok/cloud domains without explicit port, default to the scheme's default
        if configured_port is None and scheme == "https" and _needs_https(host):
            port = default_port_for_scheme
        port_part = "" if port == default_port_for_scheme else f":{port}"
        return f"{scheme}://{host}{port_part}"

    address = getattr(prompt_server, "address", "127.0.0.1") or "127.0.0.1"
    if address in ("0.0.0.0", "::"):
        address = "127.0.0.1"
    scheme = "https" if port == 443 else "http"
    default_port_for_scheme = 443 if scheme == "https" else 80
    port_part = "" if port == default_port_for_scheme else f":{port}"
    return f"{scheme}://{address}{port_part}"


def _build_worker_url(worker, endpoint=""):
    """Construct the worker base URL with optional endpoint."""
    host = (worker.get("host") or "").strip()
    port = int(worker.get("port", 8188) or 8188)

    if not host:
        host = getattr(prompt_server, "address", "127.0.0.1") or "127.0.0.1"

    if host.startswith(("http://", "https://")):
        base = host.rstrip("/")
    else:
        is_cloud = worker.get("type") == "cloud" or host.endswith(".proxy.runpod.net") or port == 443
        scheme = "https" if is_cloud else "http"
        default_port = 443 if scheme == "https" else 80
        port_part = "" if port == default_port else f":{port}"
        base = f"{scheme}://{host}{port_part}"

    return f"{base}{endpoint}"


async def _worker_is_active(worker):
    """Ping worker's /prompt endpoint to confirm it's reachable."""
    url = _build_worker_url(worker, "/prompt")
    session = await _get_client_session()
    try:
        async with session.get(url, timeout=aiohttp.ClientTimeout(total=3)) as resp:
            return resp.status == 200
    except Exception:
        return False


async def _dispatch_worker_prompt(worker, prompt_obj, workflow_meta):
    """Send the prepared prompt to a worker ComfyUI instance."""
    url = _build_worker_url(worker, "/prompt")
    payload = {"prompt": prompt_obj}
    extra_data = {}
    if workflow_meta:
        extra_data.setdefault("extra_pnginfo", {})["workflow"] = workflow_meta
    if extra_data:
        payload["extra_data"] = extra_data

    session = await _get_client_session()
    async with session.post(
        url,
        json=payload,
        timeout=aiohttp.ClientTimeout(total=60),
    ) as resp:
        resp.raise_for_status()


async def _queue_master_prompt(prompt_obj, workflow_meta, client_id):
    """Queue the master prompt through ComfyUI's prompt queue."""
    payload = {"prompt": prompt_obj}
    payload = prompt_server.trigger_on_prompt(payload)
    prompt = payload["prompt"]

    prompt_id = str(uuid.uuid4())
    valid = await execution.validate_prompt(prompt_id, prompt, None)
    if not valid[0]:
        raise RuntimeError(f"Invalid prompt: {valid[1]}")

    extra_data = {}
    if workflow_meta:
        extra_data.setdefault("extra_pnginfo", {})["workflow"] = workflow_meta
    if client_id:
        extra_data["client_id"] = client_id

    sensitive = {}
    for key in getattr(execution, "SENSITIVE_EXTRA_DATA_KEYS", []):
        if key in extra_data:
            sensitive[key] = extra_data.pop(key)

    number = getattr(prompt_server, "number", 0)
    prompt_server.number = number + 1
    prompt_queue_item = (number, prompt_id, prompt, extra_data, valid[2], sensitive)
    prompt_server.prompt_queue.put(prompt_queue_item)
    return prompt_id


def _apply_participant_overrides(
    base_prompt,
    participant_id,
    enabled_worker_ids,
    job_id_map,
    master_url,
    delegate_master,
):
    """Return a prompt copy with hidden inputs configured for master/worker."""
    prompt_copy = _deepcopy_prompt(base_prompt)
    is_master = participant_id == "master"
    worker_index_map = {wid: idx for idx, wid in enumerate(enabled_worker_ids)}
    enabled_json = json.dumps(enabled_worker_ids)

    # Distributed seed nodes
    for node_id in _find_nodes_by_class(prompt_copy, "DistributedSeed"):
        node = prompt_copy.get(node_id, {})
        inputs = node.setdefault("inputs", {})
        inputs["is_worker"] = not is_master
        if not is_master:
            idx = worker_index_map.get(participant_id, 0)
            inputs["worker_id"] = f"worker_{idx}"
        else:
            inputs["worker_id"] = ""

    # Distributed collectors
    for node_id in _find_nodes_by_class(prompt_copy, "DistributedCollector"):
        node = prompt_copy.get(node_id, {})
        inputs = node.setdefault("inputs", {})

        if _has_upstream_node(prompt_copy, node_id, "UltimateSDUpscaleDistributed"):
            inputs["pass_through"] = True
            continue

        unique_id = job_id_map.get(node_id, node_id)
        inputs["multi_job_id"] = unique_id
        inputs["is_worker"] = not is_master

        if is_master:
            inputs["enabled_worker_ids"] = enabled_json
            inputs["delegate_only"] = bool(delegate_master)
            inputs.pop("master_url", None)
            inputs.pop("worker_id", None)
        else:
            inputs["master_url"] = master_url
            inputs["worker_id"] = participant_id
            inputs["enabled_worker_ids"] = enabled_json
            inputs["delegate_only"] = False

    # Distributed upscaler nodes
    for node_id in _find_nodes_by_class(prompt_copy, "UltimateSDUpscaleDistributed"):
        node = prompt_copy.get(node_id, {})
        inputs = node.setdefault("inputs", {})
        unique_id = job_id_map.get(node_id, node_id)
        inputs["multi_job_id"] = unique_id
        inputs["is_worker"] = not is_master
        if is_master:
            inputs["enabled_worker_ids"] = enabled_json
            inputs.pop("master_url", None)
            inputs.pop("worker_id", None)
        else:
            inputs["master_url"] = master_url
            inputs["worker_id"] = participant_id
            inputs["enabled_worker_ids"] = enabled_json

    return prompt_copy


async def orchestrate_distributed_execution(
    prompt_obj,
    workflow_meta,
    client_id,
    enabled_worker_ids=None,
    delegate_master=None,
):
    """Core orchestration logic for the /distributed/queue endpoint.

    Returns:
        tuple[str, int]: (prompt_id, worker_count)
    """
    ensure_distributed_state()

    config = load_config()
    requested_ids = enabled_worker_ids if enabled_worker_ids is not None else None
    workers = _resolve_enabled_workers(config, requested_ids)

    # Respect master delegate-only configuration
    if delegate_master is None:
        delegate_master = bool(config.get("settings", {}).get("master_delegate_only", False))

    if not workers and delegate_master:
        debug_log("Delegate-only requested but no workers are enabled. Falling back to master execution.")
        delegate_master = False

    # Filter to active workers
    active_workers = []
    for worker in workers:
        if await _worker_is_active(worker):
            active_workers.append(worker)
        else:
            log(f"[Distributed] Worker {worker['name']} ({worker['id']}) is offline, skipping.")

    if not active_workers and delegate_master:
        debug_log("All workers offline while delegate-only requested; enabling master participation.")
        delegate_master = False

    enabled_ids = [worker["id"] for worker in active_workers]

    discovery_prefix = f"exec_{int(time.time() * 1000)}_{uuid.uuid4().hex[:6]}"
    job_id_map = _generate_job_id_map(prompt_obj, discovery_prefix)

    if not job_id_map:
        prompt_id = await _queue_master_prompt(prompt_obj, workflow_meta, client_id)
        return prompt_id, 0

    for job_id in job_id_map.values():
        await _ensure_distributed_queue(job_id)

    master_url = _resolve_master_url()
    master_prompt = _apply_participant_overrides(
        prompt_obj,
        "master",
        enabled_ids,
        job_id_map,
        master_url,
        delegate_master,
    )

    if delegate_master:
        collector_ids = _find_nodes_by_class(master_prompt, "DistributedCollector")
        upscale_nodes = _find_nodes_by_class(master_prompt, "UltimateSDUpscaleDistributed")
        if upscale_nodes:
            debug_log(
                "Delegate-only master mode currently does not support UltimateSDUpscaleDistributed nodes; running full prompt on master."
            )
        elif not collector_ids:
            debug_log(
                "Delegate-only master mode requested but no collectors found in master prompt. Running full prompt on master."
            )
        else:
            master_prompt = _prepare_delegate_master_prompt(master_prompt, collector_ids)

    worker_payloads = []
    for worker in active_workers:
        worker_prompt = _apply_participant_overrides(
            prompt_obj,
            worker["id"],
            enabled_ids,
            job_id_map,
            master_url,
            delegate_master,
        )
        worker_payloads.append((worker, worker_prompt))

    if worker_payloads:
        await asyncio.gather(
            *[_dispatch_worker_prompt(worker, wprompt, workflow_meta) for worker, wprompt in worker_payloads]
        )

    prompt_id = await _queue_master_prompt(master_prompt, workflow_meta, client_id)
    return prompt_id, len(worker_payloads)