from core.config import config
import contextlib
from dataclasses import dataclass
import errno
import json
import logging
import os
from pathlib import Path
from transport.legacy.port_discovery import PortDiscovery
import random
import socket
import struct
import threading
import time
from typing import Any

from models.models import MCPResponse, UnityInstanceInfo
from transport.legacy.stdio_port_registry import stdio_port_registry


logger = logging.getLogger("mcp-for-unity-server")

# Module-level lock to guard global connection initialization
_connection_lock = threading.Lock()

# Maximum allowed framed payload size (64 MiB)
FRAMED_MAX = 64 * 1024 * 1024


@dataclass
class UnityConnection:
    """Manages the socket connection to the Unity Editor."""
    host: str = config.unity_host
    port: int = None  # Will be set dynamically
    sock: socket.socket = None  # Socket for Unity communication
    use_framing: bool = False  # Negotiated per-connection
    instance_id: str | None = None  # Instance identifier for reconnection

    def __post_init__(self):
        """Set port from discovery if not explicitly provided"""
        if self.port is None:
            self.port = stdio_port_registry.get_port(self.instance_id)
        self._io_lock = threading.Lock()
        self._conn_lock = threading.Lock()
        self._needs_tool_resync = False  # Set True after reconnection

    def _prepare_socket(self, sock: socket.socket) -> None:
        try:
            sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
        except OSError as exc:
            logger.debug(f"Unable to set TCP_NODELAY: {exc}")

    def connect(self, connect_timeout: float | None = None) -> bool:
        """Establish a connection to the Unity Editor."""
        if self.sock:
            return True
        with self._conn_lock:
            if self.sock:
                return True
            try:
                # Bounded connect to avoid indefinite blocking
                if connect_timeout is None:
                    connect_timeout = float(
                        getattr(config, "connection_timeout", 1.0))
                # We trust config.unity_host (default 127.0.0.1) but future improvements
                # could dynamically prefer 'localhost' depending on OS resolver behavior.
                self.sock = socket.create_connection(
                    (self.host, self.port), connect_timeout)
                self._prepare_socket(self.sock)
                self._needs_tool_resync = True
                logger.debug(f"Connected to Unity at {self.host}:{self.port}")

                # Strict handshake: require FRAMING=1
                try:
                    require_framing = getattr(config, "require_framing", True)
                    handshake_timeout = float(
                        getattr(config, "handshake_timeout", 1.0))
                    self.sock.settimeout(handshake_timeout)
                    buf = bytearray()
                    deadline = time.monotonic() + handshake_timeout
                    while time.monotonic() < deadline and len(buf) < 512:
                        try:
                            chunk = self.sock.recv(256)
                            if not chunk:
                                break
                            buf.extend(chunk)
                            if b"\n" in buf:
                                break
                        except socket.timeout:
                            break
                    text = bytes(buf).decode('ascii', errors='ignore').strip()

                    if 'FRAMING=1' in text:
                        self.use_framing = True
                        logger.debug(
                            'MCP for Unity handshake received: FRAMING=1 (strict)')
                    else:
                        if require_framing:
                            # Best-effort plain-text advisory for legacy peers
                            with contextlib.suppress(Exception):
                                self.sock.sendall(
                                    b'MCP for Unity requires FRAMING=1\n')
                            raise ConnectionError(
                                f'MCP for Unity requires FRAMING=1, got: {text!r}')
                        else:
                            self.use_framing = False
                            logger.warning(
                                'MCP for Unity handshake missing FRAMING=1; proceeding in legacy mode by configuration')
                finally:
                    self.sock.settimeout(config.connection_timeout)
                return True
            except Exception as e:
                logger.error(f"Failed to connect to Unity: {str(e)}")
                try:
                    if self.sock:
                        self.sock.close()
                except Exception:
                    pass
                self.sock = None
                return False

    def disconnect(self):
        """Close the connection to the Unity Editor."""
        if self.sock:
            try:
                self.sock.close()
            except Exception as e:
                logger.error(f"Error disconnecting from Unity: {str(e)}")
            finally:
                self.sock = None

    def _ensure_live_connection(self) -> None:
        """Detect and discard cleanly-closed sockets before sending.

        A graceful domain reload closes the bridge's TCP connections (FIN), so a
        non-blocking peek sees EOF and lets send_command reconnect instead of
        writing to a dead socket. A half-open socket left without a FIN is not
        caught here; that case is bounded by the per-command deadline and the
        reloading-status preflight.
        """
        if not self.sock:
            return
        orig_blocking = None
        try:
            orig_blocking = self.sock.getblocking()
            self.sock.setblocking(False)
            data = self.sock.recv(1, socket.MSG_PEEK)
            if not data:
                raise ConnectionError("peer closed")
        except BlockingIOError:
            pass  # No data pending; socket is alive
        except Exception:
            logger.debug("Stale socket detected; will reconnect on next send")
            try:
                self.sock.close()
            except Exception:
                pass
            self.sock = None
        finally:
            if self.sock and orig_blocking is not None:
                self.sock.setblocking(orig_blocking)

    def _read_exact(self, sock: socket.socket, count: int) -> bytes:
        data = bytearray()
        while len(data) < count:
            chunk = sock.recv(count - len(data))
            if not chunk:
                raise ConnectionError(
                    "Connection closed before reading expected bytes")
            data.extend(chunk)
        return bytes(data)

    def receive_full_response(self, sock, buffer_size=config.buffer_size) -> bytes:
        """Receive a complete response from Unity, handling chunked data."""
        if self.use_framing:
            # Heartbeat semantics: the Unity editor emits zero-length frames while
            # a long-running command is still executing. We tolerate a bounded
            # number of these frames (or a small time window) before surfacing a
            # timeout to the caller so tools can retry or fail gracefully.
            heartbeat_limit = getattr(config, 'max_heartbeat_frames', 16)
            heartbeat_window = getattr(config, 'heartbeat_timeout', 2.0)
            heartbeat_started = time.monotonic()
            heartbeat_count = 0
            try:
                while True:
                    header = self._read_exact(sock, 8)
                    payload_len = struct.unpack('>Q', header)[0]
                    if payload_len == 0:
                        heartbeat_count += 1
                        logger.debug(
                            f"Received heartbeat frame #{heartbeat_count}")
                        if heartbeat_count >= heartbeat_limit or (time.monotonic() - heartbeat_started) > heartbeat_window:
                            raise TimeoutError(
                                "Unity sent heartbeat frames without payload within configured threshold"
                            )
                        continue
                    if payload_len > FRAMED_MAX:
                        raise ValueError(
                            f"Invalid framed length: {payload_len}")
                    payload = self._read_exact(sock, payload_len)
                    logger.debug(
                        f"Received framed response ({len(payload)} bytes)")
                    return payload
            except socket.timeout as exc:
                logger.warning("Socket timeout during framed receive")
                raise TimeoutError("Timeout receiving Unity response") from exc
            except TimeoutError:
                raise
            except Exception as exc:
                logger.error(f"Error during framed receive: {exc}")
                raise

        chunks = []
        # Respect the socket's currently configured timeout
        try:
            while True:
                chunk = sock.recv(buffer_size)
                if not chunk:
                    if not chunks:
                        raise Exception(
                            "Connection closed before receiving data")
                    break
                chunks.append(chunk)

                # Process the data received so far
                data = b''.join(chunks)
                decoded_data = data.decode('utf-8')

                # Check if we've received a complete response
                try:
                    # Special case for ping-pong
                    if decoded_data.strip().startswith('{"status":"success","result":{"message":"pong"'):
                        logger.debug("Received ping response")
                        return data

                    # Handle escaped quotes in the content
                    if '"content":' in decoded_data:
                        # Find the content field and its value
                        content_start = decoded_data.find('"content":') + 9
                        content_end = decoded_data.rfind('"', content_start)
                        if content_end > content_start:
                            # Replace escaped quotes in content with regular quotes
                            content = decoded_data[content_start:content_end]
                            content = content.replace('\\"', '"')
                            decoded_data = decoded_data[:content_start] + \
                                content + decoded_data[content_end:]

                    # Validate JSON format
                    json.loads(decoded_data)

                    # If we get here, we have valid JSON
                    logger.info(
                        f"Received complete response ({len(data)} bytes)")
                    return data
                except json.JSONDecodeError:
                    # We haven't received a complete valid JSON response yet
                    continue
                except Exception as e:
                    logger.warning(
                        f"Error processing response chunk: {str(e)}")
                    # Continue reading more chunks as this might not be the complete response
                    continue
        except socket.timeout:
            logger.warning("Socket timeout during receive")
            raise Exception("Timeout receiving Unity response")
        except Exception as e:
            logger.error(f"Error during receive: {str(e)}")
            raise

    def _cap_to_deadline(self, timeout: float, deadline: float | None, floor: float = 0.05) -> float:
        """Shrink a blocking timeout to whatever budget remains before the deadline."""
        if deadline is None:
            return timeout
        return max(floor, min(timeout, deadline - time.monotonic()))

    def send_command(self, command_type: str, params: dict[str, Any] | None = None, max_attempts: int | None = None, deadline: float | None = None) -> dict[str, Any]:
        """Send a command with retry/backoff and port rediscovery. Pings only when requested.

        Args:
            command_type: The Unity command to send
            params: Command parameters
            max_attempts: Maximum retry attempts (None = use config default, 0 = no retries)
            deadline: Shared monotonic() ceiling across retries (None = derive from command_total_timeout)
        """
        # Defensive guard: catch empty/placeholder invocations early
        if not command_type:
            raise ValueError("MCP call missing command_type")
        if params is None:
            return MCPResponse(success=False, error="MCP call received with no parameters (client placeholder?)")
        attempts = max(config.max_retries,
                       5) if max_attempts is None else max_attempts
        base_backoff = max(0.5, config.retry_delay)

        # Cap total time across all retries so a wedged socket can't block unbounded.
        total_timeout = max(0.0, float(getattr(config, "command_total_timeout", 90.0)))
        if deadline is None and total_timeout > 0:
            deadline = time.monotonic() + total_timeout

        def read_status_file(target_hash: str | None = None) -> dict | None:
            try:
                base_path = Path.home().joinpath('.unity-mcp')
                status_files = sorted(
                    base_path.glob('unity-mcp-status-*.json'),
                    key=lambda p: p.stat().st_mtime,
                    reverse=True,
                )
                if not status_files:
                    return None
                if target_hash:
                    for status_path in status_files:
                        if status_path.stem.endswith(target_hash):
                            with status_path.open('r') as f:
                                return json.load(f)
                # Fallback: return most recent regardless of hash
                with status_files[0].open('r') as f:
                    return json.load(f)
            except FileNotFoundError:
                logger.debug(
                    "Unity status file disappeared before it could be read")
                return None
            except json.JSONDecodeError as exc:
                logger.warning(f"Malformed Unity status file: {exc}")
                return None
            except OSError as exc:
                logger.warning(f"Failed to read Unity status file: {exc}")
                return None
            except Exception as exc:
                logger.debug(f"Preflight status check failed: {exc}")
                return None

        # Extract hash suffix from instance id (e.g., Project@hash)
        target_hash: str | None = None
        if self.instance_id and '@' in self.instance_id:
            maybe_hash = self.instance_id.split('@', 1)[1].strip()
            if maybe_hash:
                target_hash = maybe_hash

        # Preflight: if Unity reports reloading, return a structured hint so clients can retry politely
        try:
            status = read_status_file(target_hash)
            if status and (status.get('reloading') or status.get('reason') == 'reloading'):
                # Reload invalidates the socket; drop it under the I/O lock so this
                # close is serialized against the send/recv block, then reconnect next call.
                with self._io_lock:
                    self.disconnect()
                return MCPResponse(
                    success=False,
                    error="Unity is reloading; please retry",
                    hint="retry",
                )
        except Exception as exc:
            logger.debug(f"Preflight status check failed: {exc}")

        for attempt in range(attempts + 1):
            if deadline is not None and time.monotonic() >= deadline:
                logger.warning(
                    "Command '%s' exceeded total deadline of %.1fs after %d attempt(s); giving up",
                    command_type, total_timeout, attempt)
                raise TimeoutError(
                    f"Command '{command_type}' exceeded total deadline of "
                    f"{total_timeout:.1f}s (connection wedged or Unity unresponsive)")
            try:
                # Discard stale sockets left over from a previous domain reload
                # so we reconnect instead of writing to a dead connection.
                self._ensure_live_connection()
                # Ensure connected (handshake occurs within connect())
                t_conn_start = time.time()
                if not self.sock and not self.connect(self._cap_to_deadline(config.connection_timeout, deadline)):
                    raise ConnectionError("Could not connect to Unity")
                logger.info("[TIMING-STDIO] connect took %.3fs command=%s", time.time() - t_conn_start, command_type)

                # Build payload
                if command_type == 'ping':
                    payload = b'ping'
                else:
                    payload = json.dumps({
                        'type': command_type,
                        'params': params,
                    }).encode('utf-8')

                # Send/receive are serialized to protect the shared socket
                with self._io_lock:
                    mode = 'framed' if self.use_framing else 'legacy'
                    with contextlib.suppress(Exception):
                        logger.debug(
                            f"send {len(payload)} bytes; mode={mode}; head={payload[:32].decode('utf-8', 'ignore')}")
                    t_send_start = time.time()
                    if self.use_framing:
                        header = struct.pack('>Q', len(payload))
                        self.sock.sendall(header)
                        self.sock.sendall(payload)
                    else:
                        self.sock.sendall(payload)
                    logger.info("[TIMING-STDIO] sendall took %.3fs command=%s", time.time() - t_send_start, command_type)

                    # Cap the receive timeout to the remaining command budget (and use a
                    # short timeout during retry bursts) so a wedged socket can't block
                    # past the deadline.
                    restore_timeout = None
                    recv_timeout = 1.0 if attempt > 0 else self.sock.gettimeout()
                    if deadline is not None:
                        recv_timeout = self._cap_to_deadline(
                            recv_timeout or config.connection_timeout, deadline)
                    if recv_timeout is not None and recv_timeout != self.sock.gettimeout():
                        restore_timeout = self.sock.gettimeout()
                        self.sock.settimeout(recv_timeout)
                    try:
                        t_recv_start = time.time()
                        response_data = self.receive_full_response(self.sock)
                        logger.info("[TIMING-STDIO] receive took %.3fs command=%s len=%d", time.time() - t_recv_start, command_type, len(response_data))
                        with contextlib.suppress(Exception):
                            logger.debug(
                                f"recv {len(response_data)} bytes; mode={mode}")
                    finally:
                        if restore_timeout is not None:
                            self.sock.settimeout(restore_timeout)

                # Parse
                if command_type == 'ping':
                    resp = json.loads(response_data.decode('utf-8'))
                    if resp.get('status') == 'success' and resp.get('result', {}).get('message') == 'pong':
                        return {"message": "pong"}
                    raise Exception("Ping unsuccessful")

                resp = json.loads(response_data.decode('utf-8'))
                if resp.get('status') == 'error':
                    err = resp.get('error') or resp.get(
                        'message', 'Unknown Unity error')
                    raise Exception(err)
                return resp.get('result', {})
            except Exception as e:
                logger.warning(
                    f"Unity communication attempt {attempt+1} failed: {e}")
                try:
                    if self.sock:
                        self.sock.close()
                finally:
                    self.sock = None

                # Re-discover the port for this specific instance
                try:
                    new_port: int | None = None
                    if self.instance_id:
                        # Try to rediscover the specific instance via shared registry
                        refreshed_instance = stdio_port_registry.get_instance(
                            self.instance_id)
                        if refreshed_instance and isinstance(refreshed_instance.port, int):
                            new_port = refreshed_instance.port
                            logger.debug(
                                f"Rediscovered instance {self.instance_id} on port {new_port}")
                        else:
                            logger.warning(
                                f"Instance {self.instance_id} not found during reconnection; falling back to port scan",
                            )

                    # Fallback to registry default if instance-specific discovery failed
                    if new_port is None:
                        new_port = stdio_port_registry.get_port(
                            self.instance_id)
                        logger.info(
                            f"Using Unity port from stdio_port_registry: {new_port}")

                    if new_port != self.port:
                        logger.info(
                            f"Unity port changed {self.port} -> {new_port}")
                    self.port = new_port
                except Exception as de:
                    logger.debug(f"Port discovery failed: {de}")

                if attempt < attempts:
                    # Heartbeat-aware, jittered backoff
                    status = read_status_file(target_hash)
                    # Base exponential backoff
                    backoff = base_backoff * (2 ** attempt)
                    # Decorrelated jitter multiplier
                    jitter = random.uniform(0.1, 0.3)

                    # Fast‑retry for transient socket failures
                    fast_error = isinstance(
                        e, (ConnectionRefusedError, ConnectionResetError, TimeoutError))
                    if not fast_error:
                        try:
                            err_no = getattr(e, 'errno', None)
                            fast_error = err_no in (
                                errno.ECONNREFUSED, errno.ECONNRESET, errno.ETIMEDOUT)
                        except Exception:
                            pass

                    # Cap backoff depending on state
                    if status and status.get('reloading'):
                        # Domain reload can take 10-20s; use longer waits
                        cap = 5.0
                    elif fast_error:
                        cap = 0.25
                    else:
                        cap = 3.0

                    sleep_s = min(cap, jitter * (2 ** attempt))
                    sleep_s = self._cap_to_deadline(sleep_s, deadline, floor=0.0)
                    time.sleep(sleep_s)
                    continue
                raise


# -----------------------------
# Connection Pool for Multiple Unity Instances
# -----------------------------

class UnityConnectionPool:
    """Manages connections to multiple Unity Editor instances"""

    def __init__(self):
        self._connections: dict[str, UnityConnection] = {}
        self._known_instances: dict[str, UnityInstanceInfo] = {}
        self._last_full_scan: float = 0
        self._scan_interval: float = 5.0  # Cache for 5 seconds
        self._pool_lock = threading.Lock()
        self._default_instance_id: str | None = None

        # Check for default instance from environment
        env_default = os.environ.get("UNITY_MCP_DEFAULT_INSTANCE", "").strip()
        if env_default:
            self._default_instance_id = env_default
            logger.info(
                f"Default Unity instance set from environment: {env_default}")

    def discover_all_instances(self, force_refresh: bool = False) -> list[UnityInstanceInfo]:
        """
        Discover all running Unity Editor instances.

        Args:
            force_refresh: If True, bypass cache and scan immediately

        Returns:
            List of UnityInstanceInfo objects
        """
        now = time.time()

        # Return cached results if valid
        if not force_refresh and (now - self._last_full_scan) < self._scan_interval:
            logger.debug(
                f"Returning cached Unity instances (age: {now - self._last_full_scan:.1f}s)")
            return list(self._known_instances.values())

        # Scan for instances
        logger.debug("Scanning for Unity instances...")
        instances = PortDiscovery.discover_all_unity_instances()

        # Update cache
        with self._pool_lock:
            self._known_instances = {inst.id: inst for inst in instances}
            self._last_full_scan = now

        logger.info(
            f"Found {len(instances)} Unity instances: {[inst.id for inst in instances]}")
        return instances

    def _resolve_instance_id(self, instance_identifier: str | None, instances: list[UnityInstanceInfo]) -> UnityInstanceInfo:
        """
        Resolve an instance identifier to a specific Unity instance.

        Args:
            instance_identifier: User-provided identifier (name, hash, name@hash, path, port, or None)
            instances: List of available instances

        Returns:
            Resolved UnityInstanceInfo

        Raises:
            ConnectionError: If instance cannot be resolved
        """
        if not instances:
            raise ConnectionError(
                "No Unity Editor instances found. Please ensure Unity is running with MCP for Unity bridge."
            )

        # Use default instance if no identifier provided
        if instance_identifier is None:
            if self._default_instance_id:
                instance_identifier = self._default_instance_id
                logger.debug(f"Using default instance: {instance_identifier}")
            elif len(instances) == 1:
                # Sole instance: unambiguous, select it without requiring a hint.
                return instances[0]
            else:
                # 2+ instances connected and nothing pinned. Refuse to guess —
                # silently routing to the most-recently-heartbeated editor lets
                # an unbound session retarget another project's Unity (#1023).
                # Mirror the HTTP "multiple connected, no active set" guard.
                available_ids = [inst.id for inst in instances]
                raise ConnectionError(
                    "Multiple Unity instances are connected and none is selected. "
                    "Pass unity_instance on the call or use set_active_instance "
                    f"with one of: {available_ids}. "
                    "Read mcpforunity://instances for current sessions."
                )

        identifier = instance_identifier.strip()

        # Try exact ID match first
        for inst in instances:
            if inst.id == identifier:
                return inst

        # Try project name match
        name_matches = [inst for inst in instances if inst.name == identifier]
        if len(name_matches) == 1:
            return name_matches[0]
        elif len(name_matches) > 1:
            # Multiple projects with same name - return helpful error
            suggestions = [
                {
                    "id": inst.id,
                    "path": inst.path,
                    "port": inst.port,
                    "suggest": f"Use unity_instance='{inst.id}'"
                }
                for inst in name_matches
            ]
            raise ConnectionError(
                f"Project name '{identifier}' matches {len(name_matches)} instances. "
                f"Please use the full format (e.g., '{name_matches[0].id}'). "
                f"Available instances: {suggestions}"
            )

        # Try hash match
        hash_matches = [inst for inst in instances if inst.hash ==
                        identifier or inst.hash.startswith(identifier)]
        if len(hash_matches) == 1:
            return hash_matches[0]
        elif len(hash_matches) > 1:
            raise ConnectionError(
                f"Hash '{identifier}' matches multiple instances: {[inst.id for inst in hash_matches]}"
            )

        # Try composite format: Name@Hash or Name@Port
        if "@" in identifier:
            name_part, hint_part = identifier.split("@", 1)
            composite_matches = [
                inst for inst in instances
                if inst.name == name_part and (
                    inst.hash.startswith(hint_part) or str(
                        inst.port) == hint_part
                )
            ]
            if len(composite_matches) == 1:
                return composite_matches[0]

        # Try port match (as string)
        try:
            port_num = int(identifier)
            port_matches = [
                inst for inst in instances if inst.port == port_num]
            if len(port_matches) == 1:
                return port_matches[0]
        except ValueError:
            pass

        # Try path match
        path_matches = [inst for inst in instances if inst.path == identifier]
        if len(path_matches) == 1:
            return path_matches[0]

        # Nothing matched
        available_ids = [inst.id for inst in instances]
        raise ConnectionError(
            f"Unity instance '{identifier}' not found. "
            f"Available instances: {available_ids}. "
            f"Check mcpforunity://instances resource for all instances."
        )

    def get_connection(self, instance_identifier: str | None = None) -> UnityConnection:
        """
        Get or create a connection to a Unity instance.

        Args:
            instance_identifier: Optional identifier (name, hash, name@hash, etc.)
                                If None, uses default or most recent instance

        Returns:
            UnityConnection to the specified instance

        Raises:
            ConnectionError: If instance cannot be found or connected
        """
        # Refresh instance list if cache expired
        instances = self.discover_all_instances()

        # Resolve identifier to specific instance
        target = self._resolve_instance_id(instance_identifier, instances)

        # Return existing connection or create new one
        with self._pool_lock:
            if target.id not in self._connections:
                logger.info(
                    f"Creating new connection to Unity instance: {target.id} (port {target.port})")
                conn = UnityConnection(port=target.port, instance_id=target.id)
                if not conn.connect():
                    raise ConnectionError(
                        f"Failed to connect to Unity instance '{target.id}' on port {target.port}. "
                        f"Ensure the Unity Editor is running."
                    )
                self._connections[target.id] = conn
            else:
                # Update existing connection with instance_id and port if changed
                conn = self._connections[target.id]
                conn.instance_id = target.id
                if conn.port != target.port:
                    logger.info(
                        f"Updating cached port for {target.id}: {conn.port} -> {target.port}")
                    conn.port = target.port
                logger.debug(f"Reusing existing connection to: {target.id}")

            return self._connections[target.id]

    def disconnect_all(self):
        """Disconnect all active connections"""
        with self._pool_lock:
            for instance_id, conn in self._connections.items():
                try:
                    logger.info(
                        f"Disconnecting from Unity instance: {instance_id}")
                    conn.disconnect()
                except Exception:
                    logger.exception(f"Error disconnecting from {instance_id}")
            self._connections.clear()


# Global Unity connection pool
_unity_connection_pool: UnityConnectionPool | None = None
_pool_init_lock = threading.Lock()


def get_unity_connection_pool() -> UnityConnectionPool:
    """Get or create the global Unity connection pool"""
    global _unity_connection_pool

    if _unity_connection_pool is not None:
        return _unity_connection_pool

    with _pool_init_lock:
        if _unity_connection_pool is not None:
            return _unity_connection_pool

        logger.info("Initializing Unity connection pool")
        _unity_connection_pool = UnityConnectionPool()
        return _unity_connection_pool


# Backwards compatibility: keep old single-connection function
def get_unity_connection(instance_identifier: str | None = None) -> UnityConnection:
    """Retrieve or establish a Unity connection.

    Args:
        instance_identifier: Optional identifier for specific Unity instance.
                           If None, uses default or most recent instance.

    Returns:
        UnityConnection to the specified or default Unity instance

    Note: This function now uses the connection pool internally.
    """
    pool = get_unity_connection_pool()
    return pool.get_connection(instance_identifier)


# -----------------------------
# Centralized retry helpers
# -----------------------------

def _extract_response_reason(resp: object) -> str | None:
    """Extract a normalized (lowercase) reason string from a response.

    Returns lowercase reason values to enable case-insensitive comparisons
    by callers (e.g. _is_reloading_response, refresh_unity).
    """
    if isinstance(resp, MCPResponse):
        data = getattr(resp, "data", None)
        if isinstance(data, dict):
            reason = data.get("reason")
            if isinstance(reason, str):
                return reason.lower()
        message_text = f"{resp.message or ''} {resp.error or ''}".lower()
        if "reload" in message_text:
            return "reloading"
        return None

    if isinstance(resp, dict):
        if resp.get("state") == "reloading":
            return "reloading"
        data = resp.get("data")
        if isinstance(data, dict):
            reason = data.get("reason")
            if isinstance(reason, str):
                return reason.lower()
        message_text = (resp.get("message") or resp.get("error") or "").lower()
        if "reload" in message_text:
            return "reloading"
        return None

    return None


def _is_reloading_response(resp: object) -> bool:
    """Return True if the Unity response indicates the editor is reloading.

    Supports both raw dict payloads from Unity and MCPResponse objects returned
    by preflight checks or transport helpers.
    """
    return _extract_response_reason(resp) == "reloading"


def send_command_with_retry(
    command_type: str,
    params: dict[str, Any],
    *,
    instance_id: str | None = None,
    max_retries: int | None = None,
    retry_ms: int | None = None,
    retry_on_reload: bool = True
) -> dict[str, Any] | MCPResponse:
    """Send a command to a Unity instance, waiting politely through Unity reloads.

    Args:
        command_type: The command type to send
        params: Command parameters
        instance_id: Optional Unity instance identifier (name, hash, name@hash, etc.)
        max_retries: Maximum number of retries for reload states
        retry_ms: Delay between retries in milliseconds
        retry_on_reload: If False, don't retry when Unity is reloading (for commands
            that trigger compilation/reload and shouldn't be re-sent)

    Returns:
        Response dictionary or MCPResponse from Unity

    Uses config.reload_retry_ms and config.reload_max_retries by default. Preserves the
    structured failure if retries are exhausted.
    """
    t_retry_start = time.time()
    logger.info("[TIMING-STDIO] send_command_with_retry START command=%s", command_type)
    t_get_conn = time.time()
    conn = get_unity_connection(instance_id)
    logger.info("[TIMING-STDIO] get_unity_connection took %.3fs command=%s", time.time() - t_get_conn, command_type)
    if max_retries is None:
        max_retries = getattr(config, "reload_max_retries", 40)
    if retry_ms is None:
        retry_ms = getattr(config, "reload_retry_ms", 250)
    # Default to 20s to handle domain reloads (which can take 10-20s after tests or script changes).
    #
    # NOTE: This wait can impact agentic workflows where domain reloads happen
    # frequently (e.g., after test runs, script compilation). The 20s default
    # balances handling slow reloads vs. avoiding unnecessary delays.
    #
    # TODO: Make this more deterministic by detecting Unity's actual reload state
    # rather than blindly waiting up to 20s. See Issue #657.
    #
    # Configurable via: UNITY_MCP_RELOAD_MAX_WAIT_S (default: 20.0, max: 20.0)
    try:
        max_wait_s = float(os.environ.get(
            "UNITY_MCP_RELOAD_MAX_WAIT_S", "20.0"))
    except ValueError as e:
        raw_val = os.environ.get("UNITY_MCP_RELOAD_MAX_WAIT_S", "20.0")
        logger.warning(
            "Invalid UNITY_MCP_RELOAD_MAX_WAIT_S=%r, using default 20.0: %s",
            raw_val, e)
        max_wait_s = 20.0
    # Clamp to [0, 20] to prevent misconfiguration from causing excessive waits
    max_wait_s = max(0.0, min(max_wait_s, 20.0))

    total_timeout = max(0.0, float(getattr(config, "command_total_timeout", 90.0)))
    deadline = time.monotonic() + total_timeout if total_timeout > 0 else None

    # If retry_on_reload=False, disable connection-level retries too (issue #577)
    # Commands that trigger compilation/reload shouldn't retry on disconnect
    send_max_attempts = None if retry_on_reload else 0

    response = conn.send_command(
        command_type, params, max_attempts=send_max_attempts, deadline=deadline)
    retries = 0
    wait_started = None
    reason = _extract_response_reason(response)
    while retry_on_reload and _is_reloading_response(response) and retries < max_retries and (deadline is None or time.monotonic() < deadline):
        if wait_started is None:
            wait_started = time.monotonic()
            logger.debug(
                "Unity reload wait started: command=%s instance=%s reason=%s max_wait_s=%.2f",
                command_type,
                instance_id or "default",
                reason or "reloading",
                max_wait_s,
            )
        if max_wait_s <= 0:
            break
        elapsed = time.monotonic() - wait_started
        if elapsed >= max_wait_s:
            break
        delay_ms = retry_ms
        if isinstance(response, dict):
            retry_after = response.get("retry_after_ms")
            if retry_after is None and isinstance(response.get("data"), dict):
                retry_after = response["data"].get("retry_after_ms")
            if retry_after is not None:
                delay_ms = int(retry_after)
        sleep_ms = max(50, min(int(delay_ms), 250))
        logger.debug(
            "Unity reload wait retry: command=%s instance=%s reason=%s retry_after_ms=%s sleep_ms=%s",
            command_type,
            instance_id or "default",
            reason or "reloading",
            delay_ms,
            sleep_ms,
        )
        time.sleep(max(0.0, sleep_ms / 1000.0))
        retries += 1
        response = conn.send_command(command_type, params, deadline=deadline)
        reason = _extract_response_reason(response)

    if wait_started is not None:
        waited = time.monotonic() - wait_started
        if _is_reloading_response(response):
            logger.debug(
                "Unity reload wait exceeded budget: command=%s instance=%s waited_s=%.3f",
                command_type,
                instance_id or "default",
                waited,
            )
            return MCPResponse(
                success=False,
                error="Unity is reloading; please retry",
                hint="retry",
                data={
                    "reason": "reloading",
                    "retry_after_ms": min(250, max(50, retry_ms)),
                },
            )
        logger.debug(
            "Unity reload wait completed: command=%s instance=%s waited_s=%.3f",
            command_type,
            instance_id or "default",
            waited,
        )
    logger.info("[TIMING-STDIO] send_command_with_retry DONE total=%.3fs command=%s", time.time() - t_retry_start, command_type)
    return response


async def async_send_command_with_retry(
    command_type: str,
    params: dict[str, Any],
    *,
    instance_id: str | None = None,
    loop=None,
    max_retries: int | None = None,
    retry_ms: int | None = None,
    retry_on_reload: bool = True
) -> dict[str, Any] | MCPResponse:
    """Async wrapper that runs the blocking retry helper in a thread pool.

    Args:
        command_type: The command type to send
        params: Command parameters
        instance_id: Optional Unity instance identifier
        loop: Optional asyncio event loop
        max_retries: Maximum number of retries for reload states
        retry_ms: Delay between retries in milliseconds
        retry_on_reload: If False, don't retry when Unity is reloading

    Returns:
        Response dictionary or MCPResponse on error
    """
    try:
        import asyncio  # local import to avoid mandatory asyncio dependency for sync callers
        if loop is None:
            loop = asyncio.get_running_loop()
        result = await loop.run_in_executor(
            None,
            lambda: send_command_with_retry(
                command_type, params, instance_id=instance_id, max_retries=max_retries,
                retry_ms=retry_ms, retry_on_reload=retry_on_reload),
        )

        # After a successful command, check if the connection was freshly
        # established (reconnection after domain reload).  If so, re-sync
        # tool visibility and custom tool registration from Unity.
        # Always clear the flag, but only schedule the background resync
        # when this call is not itself get_tool_states (to avoid recursion).
        try:
            pool = get_unity_connection_pool()
            conn = pool.get_connection(instance_id)
            if getattr(conn, "_needs_tool_resync", False):
                conn._needs_tool_resync = False
                if command_type != "get_tool_states":
                    logger.info(
                        "Detected reconnection to Unity; scheduling tool re-sync"
                    )
                    asyncio.ensure_future(_resync_tools_after_reconnect(instance_id))
        except Exception as exc:
            logger.debug(
                "Failed to schedule post-reconnection tool re-sync: %s",
                exc,
            )

        return result
    except Exception as e:
        return MCPResponse(success=False, error=str(e))


async def _resync_tools_after_reconnect(instance_id: str | None) -> None:
    """Background task: re-sync tool visibility and custom tools after reconnection."""
    try:
        from services.tools import sync_tool_visibility_from_unity
        result = await sync_tool_visibility_from_unity(
            instance_id=instance_id, notify=True,
        )
        if result.get("synced"):
            logger.info(
                "Post-reconnection tool re-sync complete: "
                "enabled=[%s], disabled=[%s], custom_tools=%d",
                ", ".join(result.get("enabled_groups", [])),
                ", ".join(result.get("disabled_groups", [])),
                result.get("custom_tool_count", 0),
            )
        else:
            logger.debug(
                "Post-reconnection tool re-sync skipped: %s",
                result.get("error", "unknown"),
            )
    except Exception as exc:
        logger.debug("Post-reconnection tool re-sync failed: %s", exc)
