"""cssl_sharing.py - CSSL Runtime Sharing / Inter-Runtime Communication

v4.9.15: Builtin include for connecting CSSL runtimes across threads or processes.

Usage in CSSL:
    auto cshare = include("cssl-sharing");
    cshare = new CsslSharing(channel="my_channel", portID=1234);
    cshare.open();                         // Register in channel
    cshare.validate(port=5678);            // Whitelist a port (optional)
    cshare.connect(port=5678);            // Connect to another runtime
    cshare.interact(port=5678, FUNC);     // Execute function on target
    cshare.merge(port=5678);             // Merge target's scope into current
    info = cshare.getRuntimeInfo();      // Dict of current runtime state

Usage in Python:
    from includecpp.core.cssl_bridge import CsslLang
    cssl = CsslLang()
    cssl_sharing = cssl.Sharing(channel="my_channel", portID=12345)
    cssl_sharing.open()
    cssl_sharing.connect(port=target_port)
    cssl_sharing.interact(port=target_port, func=lambda rt: rt.run('printl("hi")'))
"""
from __future__ import annotations

import json
import socket
import threading
import time
import uuid
from typing import Any, Callable, Dict, List, Optional, Set


# ─────────────────────────────────────────────────────────────────────────────
# Global in-process channel registry
# All CsslSharing instances in the same Python process share this dict.
# ─────────────────────────────────────────────────────────────────────────────

_CHANNELS: Dict[str, Dict[int, 'CsslSharing']] = {}
_CHANNELS_LOCK = threading.Lock()


class CSSLSharingError(Exception):
    """Raised when a sharing operation fails."""


# ─────────────────────────────────────────────────────────────────────────────
# Socket protocol helpers (cross-process communication)
# ─────────────────────────────────────────────────────────────────────────────

_ENCODING = 'utf-8'
_HEADER_SIZE = 8          # 8-byte big-endian length prefix per message
_SOCKET_TIMEOUT = 10.0   # seconds


def _send_msg(conn: socket.socket, payload: dict) -> None:
    """Send a JSON payload with a length-prefix header."""
    data = json.dumps(payload, default=str).encode(_ENCODING)
    header = len(data).to_bytes(_HEADER_SIZE, 'big')
    conn.sendall(header + data)


def _recv_msg(conn: socket.socket) -> dict:
    """Receive a length-prefixed JSON payload."""
    header = _recv_exact(conn, _HEADER_SIZE)
    length = int.from_bytes(header, 'big')
    data = _recv_exact(conn, length)
    return json.loads(data.decode(_ENCODING))


def _recv_exact(conn: socket.socket, n: int) -> bytes:
    buf = bytearray()
    while len(buf) < n:
        chunk = conn.recv(n - len(buf))
        if not chunk:
            raise CSSLSharingError("Connection closed by peer")
        buf.extend(chunk)
    return bytes(buf)


def _call_remote(host: str, port: int, payload: dict, timeout: float = _SOCKET_TIMEOUT) -> dict:
    """Send a request to a remote CsslSharing socket server and return response."""
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
        s.settimeout(timeout)
        s.connect((host, port))
        _send_msg(s, payload)
        return _recv_msg(s)


# ─────────────────────────────────────────────────────────────────────────────
# Peer descriptor
# ─────────────────────────────────────────────────────────────────────────────

class _SharingPeer:
    """Represents a connected peer — either in-process or remote."""

    def __init__(
        self,
        port: int,
        channel: str,
        sharing_obj: Optional['CsslSharing'] = None,  # in-process
        host: str = '127.0.0.1',
    ) -> None:
        self.port = port
        self.channel = channel
        self._sharing = sharing_obj   # in-process reference
        self._host = host             # for remote peers

    @property
    def is_local(self) -> bool:
        return self._sharing is not None

    @property
    def runtime(self) -> Any:
        """Return the underlying CSSLRuntime, resolving via bridge if needed."""
        if self._sharing is None:
            return None
        if self._sharing._bridge is not None:
            return self._sharing._bridge._get_runtime()
        return self._sharing._runtime

    def execute_code(self, code: str) -> Any:
        """Execute CSSL code on the peer's runtime and return result."""
        if self.is_local:
            # Prefer bridge path — no direct internal class calls
            if self._sharing._bridge is not None:
                return self._sharing._bridge.run(code)
            # Fallback: direct runtime (when called from within a CSSL builtin)
            rt = self._sharing._runtime
            if rt is None:
                raise CSSLSharingError(f"Peer {self.port} has no active runtime")
            from .cssl_parser import parse_cssl_program
            from .cssl_runtime import CSSLReturn
            ast = parse_cssl_program(code)
            result = None
            for node in (ast.children if ast else []):
                try:
                    result = rt._execute_node(node)
                except CSSLReturn as ret:
                    return ret.value
            return result
        else:
            resp = _call_remote(self._host, self.port, {'cmd': 'execute', 'code': code})
            if resp.get('status') == 'error':
                raise CSSLSharingError(f"Remote execute error: {resp.get('msg')}")
            return resp.get('data')

    def get_scope_snapshot(self) -> dict:
        """Return a JSON-safe snapshot of the peer's global scope."""
        if self.is_local:
            rt = self.runtime  # resolves via bridge if available
            if rt is None:
                return {}
            raw = {}
            scope = rt.global_scope
            if hasattr(scope, '_vars'):
                for name, val in scope._vars.items():
                    if not name.startswith('_'):
                        try:
                            json.dumps(val, default=str)  # only JSON-safe values
                            raw[name] = val
                        except Exception:
                            pass
            return raw
        else:
            resp = _call_remote(self._host, self.port, {'cmd': 'get_scope'})
            if resp.get('status') == 'error':
                raise CSSLSharingError(f"Remote get_scope error: {resp.get('msg')}")
            return resp.get('data', {})

    def find(self, kind: str, name: str) -> Any:
        """Find a function or variable on the peer's runtime."""
        if self.is_local:
            rt = self.runtime  # resolves via bridge if available
            if rt is None:
                return None
            val = rt.scope.get(name) or rt.global_scope.get(name)
            if kind == 'function':
                if val and hasattr(val, 'type') and val.type == 'function':
                    return val
                if hasattr(rt, 'builtins') and rt.builtins.has_function(name):
                    return rt.builtins._functions.get(name)
                return None
            return val  # variable
        else:
            resp = _call_remote(self._host, self.port,
                                {'cmd': 'find', 'kind': kind, 'name': name})
            if resp.get('status') == 'error':
                return None
            return resp.get('data')

    def ping(self) -> bool:
        if self.is_local:
            return True
        try:
            resp = _call_remote(self._host, self.port, {'cmd': 'ping'}, timeout=3.0)
            return resp.get('status') == 'ok'
        except Exception:
            return False

    def get_runtime_info(self) -> dict:
        if self.is_local:
            return self._sharing.getRuntimeInfo()
        resp = _call_remote(self._host, self.port, {'cmd': 'info'})
        return resp.get('data', {})


# ─────────────────────────────────────────────────────────────────────────────
# Main CsslSharing class
# ─────────────────────────────────────────────────────────────────────────────

class CsslSharing:
    """CSSL Runtime Sharing — bidirectional inter-runtime communication.

    Connects two or more CSSL runtimes (same process or different processes)
    through a named channel and numeric port IDs.

    Constructor parameters:
        channel (str)  : Shared channel name — all runtimes using the same
                         channel name can discover each other.
        portID  (int)  : Unique port number that identifies this runtime in the
                         channel. For cross-process communication the socket
                         server listens on this TCP port (localhost).

    Usage in CSSL:
        auto cshare = include("cssl-sharing");
        cshare = new CsslSharing(channel="game_sync", portID=4000);
        cshare.open();
        cshare.validate(port=4001);    // Only accept connections from 4001
        cshare.connect(port=4001);
        cshare.interact(port=4001, executeCode("printl('hello from other side')"));
        result = cshare.find(port=4001, type="variable", name="score");
        cshare.merge(port=4001);
        cshare.close();

    Usage in Python:
        sharing = cssl.Sharing(channel="game_sync", portID=4000)
        sharing.open()
        sharing.connect(port=4001)
        sharing.interact(port=4001, lambda rt: rt.run('x = x + 1'))
    """

    def __init__(
        self,
        channel: str = 'default',
        portID: int = 0,
    ) -> None:
        self.channel: str = str(channel)
        self.portID: int = int(portID)
        self._runtime: Any = None           # Set by _get_sharing_module() (CSSLRuntime)
        self._bridge: Any = None            # Set by CsslLang.Sharing() (CsslLang bridge)
        self._peers: Dict[int, _SharingPeer] = {}
        self._whitelist: Set[int] = set()
        self._is_open: bool = False
        self._server_sock: Optional[socket.socket] = None
        self._server_thread: Optional[threading.Thread] = None
        self._lock = threading.Lock()

    # ── Lifecycle ────────────────────────────────────────────────────────────

    def open(self) -> 'CsslSharing':
        """Register this runtime in the channel and start the socket server.

        After open(), other runtimes in the same channel can connect to this
        one by portID. Cross-process connections come in via TCP on portID.
        """
        with _CHANNELS_LOCK:
            if self.channel not in _CHANNELS:
                _CHANNELS[self.channel] = {}
            _CHANNELS[self.channel][self.portID] = self

        self._start_server()
        self._is_open = True
        return self

    def close(self) -> None:
        """Unregister from the channel and close all connections."""
        with _CHANNELS_LOCK:
            channel = _CHANNELS.get(self.channel, {})
            channel.pop(self.portID, None)
            if not channel:
                _CHANNELS.pop(self.channel, None)

        if self._server_sock:
            try:
                self._server_sock.close()
            except Exception:
                pass
            self._server_sock = None

        with self._lock:
            self._peers.clear()
        self._is_open = False

    # ── Access control ───────────────────────────────────────────────────────

    def validate(self, port: int) -> 'CsslSharing':
        """Restrict incoming connect() calls to the given portID only.

        Must be called before connect() on the other side. Any port not in
        the whitelist will be rejected.

        Args:
            port: The portID that is allowed to connect.
        """
        self._whitelist.add(int(port))
        return self

    # ── Connection ───────────────────────────────────────────────────────────

    def connect(self, port: int) -> 'CsslSharing':
        """Connect to another runtime by portID in the same channel.

        First tries in-process lookup (same Python process). Falls back to
        TCP socket connection on localhost:port.

        Args:
            port: portID of the target runtime.

        Raises:
            CSSLSharingError: If the target whitelist blocks this connection.
        """
        port = int(port)

        # In-process lookup
        with _CHANNELS_LOCK:
            channel = _CHANNELS.get(self.channel, {})
            peer_sharing = channel.get(port)

        if peer_sharing is not None:
            # Check target's whitelist
            if peer_sharing._whitelist and self.portID not in peer_sharing._whitelist:
                raise CSSLSharingError(
                    f"Port {self.portID} rejected by port {port}'s whitelist"
                )
            peer = _SharingPeer(port, self.channel, sharing_obj=peer_sharing)
        else:
            # Cross-process: try TCP
            peer = _SharingPeer(port, self.channel, host='127.0.0.1')
            resp = _call_remote('127.0.0.1', port, {
                'cmd': 'handshake',
                'from_port': self.portID,
                'channel': self.channel,
            })
            if resp.get('status') == 'error':
                raise CSSLSharingError(
                    f"Connection to port {port} refused: {resp.get('msg')}"
                )

        with self._lock:
            self._peers[port] = peer
        return self

    # ── Data / control ───────────────────────────────────────────────────────

    def get(self, port: int) -> Optional[_SharingPeer]:
        """Return the _SharingPeer for a connected port, or null if not connected."""
        return self._peers.get(int(port))

    def exists(self, port: int) -> bool:
        """Return True if a connection to the given port is active."""
        port = int(port)
        peer = self._peers.get(port)
        if peer is None:
            return False
        return peer.ping()

    def getRuntimeInfo(self, key: str = None) -> Any:
        """Return a dict of the current runtime's state.

        Args:
            key: Optional string key to extract a single field.
                 Values: 'status', 'port', 'channel', 'peers',
                         'variables', 'functions'

        Returns:
            Dict if no key, or the specific value if key is given.
        """
        rt = self._bridge._get_runtime() if self._bridge is not None else self._runtime
        vars_list: List[str] = []
        funcs_list: List[str] = []
        if rt is not None:
            for scope in (rt.scope, rt.global_scope):
                if hasattr(scope, '_vars'):
                    for name, val in scope._vars.items():
                        if name.startswith('_'):
                            continue
                        if hasattr(val, 'type') and val.type == 'function':
                            if name not in funcs_list:
                                funcs_list.append(name)
                        elif name not in vars_list:
                            vars_list.append(name)

        info: Dict[str, Any] = {
            'status': 'connected' if self._is_open else 'disconnected',
            'port': self.portID,
            'channel': self.channel,
            'peers': list(self._peers.keys()),
            'variables': vars_list,
            'functions': funcs_list,
        }
        if key is not None:
            return info.get(str(key))
        return info

    def interact(self, port: int, func: Any) -> Any:
        """Execute a callable or CSSL code on the target runtime.

        Args:
            port: Connected portID of the target runtime.
            func: One of:
                  - Python callable: called with (target_CsslSharing_obj)
                  - CSSL function ASTNode: executed on the target runtime
                  - str: CSSL code string executed on the target runtime

        Returns:
            The return value of the function / last expression.
        """
        port = int(port)
        peer = self._peers.get(port)
        if peer is None:
            raise CSSLSharingError(f"No active connection to port {port}")

        if callable(func):
            # Python callable: pass target bridge (preferred) or sharing object
            if peer.is_local:
                target = peer._sharing._bridge if peer._sharing._bridge is not None else peer._sharing
                return func(target)
            raise CSSLSharingError(
                "Python callables can only be used for same-process interact()"
            )

        if isinstance(func, str):
            return peer.execute_code(func)

        # CSSL function ASTNode — use execute_code to avoid direct runtime calls
        if hasattr(func, 'type') and func.type == 'function':
            if peer.is_local:
                name = func.value.get('name', '') if isinstance(func.value, dict) else ''
                if name:
                    return peer.execute_code(f'{name}();')
                # Fallback: direct runtime call for anonymous functions
                rt = peer.runtime
                if rt is not None and hasattr(rt, '_call_function_by_node'):
                    return rt._call_function_by_node(func, [], {})
            # Serialize: get source (if stored) or fall back to name
            name = func.value.get('name', '') if isinstance(func.value, dict) else ''
            if name:
                return peer.execute_code(f'{name}();')
            raise CSSLSharingError(
                "Cannot serialize anonymous CSSL function for cross-process interact()"
            )

        raise CSSLSharingError(
            f"interact() func must be a callable, str, or CSSL function ASTNode, got {type(func).__name__}"
        )

    def merge(self, port: int) -> None:
        """Merge the target runtime's public scope into this runtime.

        All public variables and functions from the target are copied into the
        current runtime's global scope. Changes made in either runtime after
        merging are NOT automatically synchronized (snapshot at merge time).

        Args:
            port: Connected portID of the target runtime.
        """
        port = int(port)
        peer = self._peers.get(port)
        if peer is None:
            raise CSSLSharingError(f"No active connection to port {port}")

        self_rt = self._bridge._get_runtime() if self._bridge is not None else self._runtime
        peer_rt = peer.runtime  # resolves via peer's bridge or runtime
        if peer.is_local and peer_rt is not None and self_rt is not None:
            target_scope = peer_rt.global_scope
            if hasattr(target_scope, '_vars'):
                for name, val in list(target_scope._vars.items()):
                    if not name.startswith('_'):
                        if self._bridge is not None:
                            self._bridge.set_global(name, val)
                        else:
                            self_rt.global_scope.set(name, val)
        else:
            # Cross-process: fetch JSON-safe snapshot and apply
            snapshot = peer.get_scope_snapshot()
            for name, val in snapshot.items():
                if self._bridge is not None:
                    self._bridge.set_global(name, val)
                elif self_rt is not None:
                    self_rt.global_scope.set(name, val)

    def find(self, port: int, type: str = 'function', name: str = '') -> Any:
        """Find a named function or variable on a connected runtime.

        Args:
            port: Connected portID.
            type: 'function' or 'variable'
            name: The name to look up.

        Returns:
            The function / variable value, or null if not found.
        """
        port = int(port)
        peer = self._peers.get(port)
        if peer is None:
            return None
        return peer.find(str(type), str(name))

    # ── Socket server (cross-process) ────────────────────────────────────────

    def _start_server(self) -> None:
        """Start the TCP server thread that listens on localhost:portID."""
        if self.portID == 0:
            return  # No socket server for port 0 (in-process only)
        try:
            srv = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            srv.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
            srv.bind(('127.0.0.1', self.portID))
            srv.listen(16)
            srv.settimeout(1.0)
            self._server_sock = srv

            t = threading.Thread(
                target=self._server_loop,
                name=f'CsslSharing-{self.channel}-{self.portID}',
                daemon=True,
            )
            t.start()
            self._server_thread = t
        except OSError:
            # Port already in use or unavailable — in-process only mode
            self._server_sock = None

    def _server_loop(self) -> None:
        srv = self._server_sock
        while srv and self._is_open:
            try:
                conn, addr = srv.accept()
            except socket.timeout:
                continue
            except Exception:
                break
            t = threading.Thread(
                target=self._handle_client,
                args=(conn, addr),
                daemon=True,
            )
            t.start()

    def _handle_client(self, conn: socket.socket, addr: tuple) -> None:
        with conn:
            try:
                msg = _recv_msg(conn)
                resp = self._dispatch_server_cmd(msg)
                _send_msg(conn, resp)
            except Exception as exc:
                try:
                    _send_msg(conn, {'status': 'error', 'msg': str(exc)})
                except Exception:
                    pass

    def _dispatch_server_cmd(self, msg: dict) -> dict:
        cmd = msg.get('cmd', '')

        if cmd == 'ping':
            return {'status': 'ok'}

        if cmd == 'handshake':
            from_port = int(msg.get('from_port', 0))
            if self._whitelist and from_port not in self._whitelist:
                return {'status': 'error', 'msg': f'Port {from_port} not in whitelist'}
            return {'status': 'ok', 'port': self.portID, 'channel': self.channel}

        if cmd == 'execute':
            code = msg.get('code', '')
            try:
                if self._runtime is None:
                    return {'status': 'error', 'msg': 'Runtime not available'}
                from .cssl_parser import parse_cssl_program
                from .cssl_runtime import CSSLReturn
                ast = parse_cssl_program(code)
                result = None
                for node in (ast.children if ast else []):
                    try:
                        result = self._runtime._execute_node(node)
                    except CSSLReturn as ret:
                        result = ret.value
                        break
                return {'status': 'ok', 'data': result}
            except Exception as exc:
                return {'status': 'error', 'msg': str(exc)}

        if cmd == 'get_scope':
            if self._runtime is None:
                return {'status': 'ok', 'data': {}}
            raw: dict = {}
            for scope in (self._runtime.scope, self._runtime.global_scope):
                if hasattr(scope, '_vars'):
                    for name, val in scope._vars.items():
                        if name.startswith('_') or name in raw:
                            continue
                        try:
                            json.dumps(val, default=str)
                            raw[name] = val
                        except Exception:
                            pass
            return {'status': 'ok', 'data': raw}

        if cmd == 'find':
            kind = msg.get('kind', 'function')
            name = msg.get('name', '')
            if self._runtime is None:
                return {'status': 'ok', 'data': None}
            val = self._runtime.scope.get(name) or self._runtime.global_scope.get(name)
            if kind == 'function':
                if val and hasattr(val, 'type') and val.type == 'function':
                    return {'status': 'ok', 'data': {'name': name, 'type': 'function'}}
                return {'status': 'ok', 'data': None}
            try:
                serialized = json.dumps(val, default=str)
                return {'status': 'ok', 'data': json.loads(serialized)}
            except Exception:
                return {'status': 'ok', 'data': str(val)}

        if cmd == 'info':
            return {'status': 'ok', 'data': self.getRuntimeInfo()}

        return {'status': 'error', 'msg': f'Unknown command: {cmd}'}


# ─────────────────────────────────────────────────────────────────────────────
# Module object — returned by include("cssl-sharing")
# ─────────────────────────────────────────────────────────────────────────────

class CsslSharingModule:
    """CSSL Sharing module — returned by include("cssl-sharing").

    Usage:
        auto cshare = include("cssl-sharing");
        cshare = new CsslSharing(channel="sync", portID=4000);
        cshare.open();
        cshare.connect(port=4001);
    """

    CsslSharing = CsslSharing

    def __init__(self) -> None:
        self._runtime: Any = None

    def __getattr__(self, name: str) -> Any:
        if name == 'CsslSharing':
            return CsslSharing
        raise AttributeError(f"'cssl-sharing' module has no attribute '{name}'")


_sharing_module_instance: Optional[CsslSharingModule] = None


def get_sharing_module() -> CsslSharingModule:
    global _sharing_module_instance
    if _sharing_module_instance is None:
        _sharing_module_instance = CsslSharingModule()
    return _sharing_module_instance
