"""
TCP Server for Electron ↔ Python IPC (Linux / macOS).

On Windows the equivalent is pipe_server.py (Named Pipes).

Wire protocol
-------------
All messages are newline-delimited JSON sent over a localhost TCP connection
to 127.0.0.1:54321.  Each message is a single JSON object followed by '\\n'.

Message format (both directions)::

    {"type": "<message_type>", "data": {…}}

Handshake sequence (must complete within HANDSHAKE_TIMEOUT seconds):
  1. Electron  → Python : {"type": "hello_from_electron"}
  2. Python    → Electron: {"type": "hello_from_python", "timestamp": <float>}
  3. Python    → Electron: {"type": "ready", "data": {"status": "ready",
                             "capabilities": [...], "version": "..."}}

After handshake any JSON message may flow in either direction.

Common message types sent BY Python
------------------------------------
  ready              — sent after handshake; announces capabilities.
  log                — {"level": "info|warning|error", "message": "..."}
  screenshot_captured— empty data; confirms a screenshot was just saved.
  tracking_error     — {"code": "[Exxx]", "message": "...", "recoverable": bool}
                       Surface to user when tracking silently breaks.
  monitoring_status  — {"is_running": bool, ...}
  auth_failed        — {"reason": "..."}
  stop_completed     — {"status": "success|error"}
  timezone_error     — {"message": "..."}

Common message types sent BY Electron
---------------------------------------
  hello_from_electron— handshake initiation.
  start              — begin tracking; {"projectId": ..., "taskId": ...}
  stop               — stop tracking.
  set_project_task   — {"projectId": ..., "taskId": ...}
  get_monitoring_status — request current status.
  idle_popup_response— user response to idle popup.

Error codes
-----------
See ws_package/utils/error_codes.py for the full [Exxx] catalogue.
[E105] IPC server failed to start.
[E106] TCP port 54321 already in use (antivirus / stale process).
"""
import socket
import select
import json
import threading
import queue
import time
import struct
import os
import uuid
from typing import Optional, Dict, Any
from ws_package.log.writer import LogWriter
from ws_package.utils.error_codes import AppError, ErrorCode

logger = LogWriter.file(__name__)

# Default TCP configuration
TCP_HOST = '127.0.0.1'  # Use 127.0.0.1, NOT localhost
TCP_PORT = 54321
TCP_BUFFER_SIZE = 4096
TCP_TIMEOUT = 5.0  # Socket timeout in seconds
HANDSHAKE_TIMEOUT = 5.0  # 5 seconds to complete handshake


class TCPServer:
    """TCP server for bidirectional communication with Electron"""
    
    def __init__(self, host: str = TCP_HOST, port: int = TCP_PORT):
        self.host = host
        self.port = port
        self.server_socket: Optional[socket.socket] = None
        self.client_socket: Optional[socket.socket] = None
        self.client_address: Optional[tuple] = None
        self.client_handshake_complete = False  # Track handshake completion
        self.session_id = str(uuid.uuid4())[:8]  # Short session ID for log correlation
        self.running = False
        self.message_queue = queue.Queue(maxsize=1000)
        self.receive_thread: Optional[threading.Thread] = None
        self.send_lock = threading.Lock()
        self.client_lock = threading.Lock()  # Lock for client_socket replacement
        self.receive_buffer = b''
        
    def _check_and_cleanup_port(self):
        """Check if port is in use and try to clean it up"""
        try:
            import psutil
            current_pid = os.getpid()
            
            for conn in psutil.net_connections(kind='inet'):
                if conn.status in ('LISTEN', 'TIME_WAIT') and conn.laddr.port == self.port:
                    if conn.pid == current_pid:
                        logger.info("Port %s is bound by this process (PID %s) - will attempt to reuse", self.port, current_pid)
                        return True
                    else:
                        logger.warning("Port %s is in use by process %s (current PID: %s)", self.port, conn.pid, current_pid)
                        # Try to kill the process if it's a Python process (be careful!)
                        try:
                            proc = psutil.Process(conn.pid)
                            if 'python' in proc.name().lower() or 'workstatus' in ' '.join(proc.cmdline()).lower():
                                logger.warning("Attempting to terminate process %s using port %s...", conn.pid, self.port)
                                proc.terminate()
                                proc.wait(timeout=2)
                                logger.info("Process %s terminated successfully", conn.pid)
                                time.sleep(0.5)  # Wait for port to be released
                                return True
                        except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.TimeoutExpired):
                            logger.warning("Could not terminate process %s", conn.pid)
            return False
        except ImportError:
            # psutil not available
            return False
        except Exception as e:
            logger.debug("Error checking port: %s", e)
            return False
    
    def start(self):
        """Start the TCP server and listen for connections"""
        if self.running:
            logger.warning("TCP server already running")
            return
        
        # Check if port is in use and try to clean it up
        self._check_and_cleanup_port()
        
        try:
            # Bind with short retry loop.  SO_REUSEADDR handles TIME_WAIT
            # states instantly; real port conflicts (stale process) are
            # resolved by _check_and_cleanup_port() before we enter this loop.
            max_retries = 5
            retry_delay = 0.1  # start at 100 ms, doubles each attempt (max 1 s)

            for attempt in range(max_retries):
                try:
                    self.server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
                    self.server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
                    try:
                        self.server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEPORT, 1)
                    except (AttributeError, OSError):
                        pass
                    try:
                        self.server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_LINGER, struct.pack('ii', 1, 0))
                    except (AttributeError, OSError):
                        pass
                    self.server_socket.settimeout(TCP_TIMEOUT)
                    self.server_socket.bind((self.host, self.port))
                    break  # success

                except OSError as e:
                    if e.errno in (98, 48):  # EADDRINUSE (Linux / macOS)
                        try:
                            if self.server_socket:
                                self.server_socket.close()
                        except Exception:
                            pass
                        self.server_socket = None

                        if attempt == 0:
                            # First failure — try killing the stale occupant
                            self._check_and_cleanup_port()

                        if attempt < max_retries - 1:
                            logger.warning("%s — port %s in use (attempt %s/%s), retrying in %.1fs...",
                                           ErrorCode.IPC_PORT_IN_USE, self.port, attempt + 1, max_retries, retry_delay)
                            time.sleep(retry_delay)
                            retry_delay = min(retry_delay * 2, 1.0)
                            continue
                        else:
                            logger.critical("%s — port %s still in use after %s attempts. "
                                            "Check for stale Workstatus process or antivirus port blocking.",
                                            ErrorCode.IPC_PORT_IN_USE, self.port, max_retries)
                            raise AppError(ErrorCode.IPC_PORT_IN_USE,
                                           f"{ErrorCode.IPC_PORT_IN_USE} — port {self.port} still in use after "
                                           f"{max_retries} attempts. Ensure no other instance is running.",
                                           is_recoverable=False,
                                           recovery_hint="Check for stale Workstatus process or antivirus port blocking.")
                    else:
                        try:
                            if self.server_socket:
                                self.server_socket.close()
                        except Exception:
                            pass
                        self.server_socket = None
                        raise
            
            self.server_socket.listen(1)  # Allow only one connection at a time
            self.running = True
            
            # Print to stdout immediately (logger might not be ready)
            print(f"[Timing][Python] TCP server bound to {self.host}:{self.port}", flush=True)
            logger.info("TCP server started on %s:%s", self.host, self.port)
            
            # Start accept thread
            accept_thread = threading.Thread(
                target=self._accept_connections,
                daemon=True,
                name="TCPAcceptThread"
            )
            accept_thread.start()
            
            # Print confirmation
            print(f"[Timing][Python] TCP server listening and ready for connections", flush=True)
            
        except OSError as e:
            if e.errno in (98, 48):
                logger.critical("%s — port %s already in use: %s",
                                ErrorCode.IPC_PORT_IN_USE, self.port, e)
                raise AppError(ErrorCode.IPC_PORT_IN_USE,
                               f"{ErrorCode.IPC_PORT_IN_USE} — port {self.port} is already in use.",
                               is_recoverable=False,
                               recovery_hint="Ensure no other Workstatus instance is running and antivirus is not blocking port binding.")
            else:
                logger.critical("%s: %s", ErrorCode.IPC_SERVER_FAILED, e, exc_info=True)
                raise
        except Exception as e:
            logger.critical("%s: %s", ErrorCode.IPC_SERVER_FAILED, e, exc_info=True)
            self.running = False
            raise
    
    def _accept_connections(self):
        """Accept incoming connections from Electron clients"""
        logger.info("Waiting for Electron client connection...")
        last_log_time = time.time()
        while self.running:
            try:
                readable, _, _ = select.select([self.server_socket], [], [], 0.5)
                if not readable:
                    # Only log "still waiting" if we don't have an active client connection
                    # Check if client is connected and handshake is complete
                    with self.client_lock:
                        has_active_client = (self.client_socket is not None and 
                                          self.client_handshake_complete)
                    
                    if not has_active_client:
                        # Log every 30 seconds that we're still waiting (only if no active client)
                        current_time = time.time()
                        if current_time - last_log_time >= 30.0:
                            logger.info("Still waiting for Electron client connection... (server running: %s, port: %s)", self.running, self.port)
                            last_log_time = current_time
                    # If we have an active client, don't log - just continue silently
                    continue
                    
                logger.info("Incoming connection detected, accepting...")
                client_sock, client_addr = self.server_socket.accept()
                logger.info("Connection accepted from %s, setting up...", client_addr)
                
                # Capture and evict any existing client WITHOUT holding client_lock
                # while calling _disconnect_client(), avoiding the deadlock that
                # occurs when _disconnect_client() tries to re-acquire client_lock.
                old_socket = None
                old_thread = None
                with self.client_lock:
                    if self.client_socket:
                        old_addr = self.client_address
                        old_socket = self.client_socket
                        old_thread = self.receive_thread
                        # Clear state inline (same as _disconnect_client but without re-acquiring the lock)
                        self.client_socket = None
                        self.client_address = None
                        self.client_handshake_complete = False
                        self.receive_buffer = b''
                        logger.info("Replacing existing connection from %s with new connection from %s", old_addr, client_addr)
                        
                        # Clean up old client state first (sets self.client_socket to None)
                        self._disconnect_client()
                        
                        # Wait for old receive thread to finish before accepting new connection
                        if old_thread and old_thread.is_alive():
                            logger.info("Waiting for old receive thread to finish...")
                            # Close old socket to signal thread to exit (recv will fail)
                            try:
                                old_socket.shutdown(socket.SHUT_RDWR)
                            except Exception:
                                pass
                            try:
                                old_socket.close()
                            except Exception:
                                pass
                            # Wait for thread to finish (with timeout)
                            old_thread.join(timeout=2.0)
                            if old_thread.is_alive():
                                logger.warning("Old receive thread did not finish in time")
                    
                    # Accept the new client
                    self.client_socket = client_sock
                    self.client_address = client_addr
                    # Set a longer timeout for receive operations (30 seconds) to avoid premature timeouts
                    self.client_socket.settimeout(30.0)
                    self.receive_buffer = b''
                    self.client_handshake_complete = False  # Reset handshake status
                    logger.info("New client socket configured: timeout=30.0s, address=%s", client_addr)
                
                logger.info("Electron client connected from %s", client_addr)
                
                # Small delay to ensure socket is fully ready before starting receive thread
                time.sleep(0.05)
                
                # Start receive thread for this client (handshake will be handled in receive thread)
                with self.client_lock:
                    current_socket = self.client_socket
                    if current_socket:
                        # Verify socket is still the same one we just accepted
                        if current_socket is not client_sock:
                            logger.warning("Socket mismatch! Expected %s, got %s", client_sock, current_socket)
                        logger.info("Starting receive thread for socket from %s...", client_addr)
                        self.receive_thread = threading.Thread(
                            target=self._receive_messages_with_socket,
                            args=(current_socket,),
                            daemon=True,
                            name=f"TCPReceiveThread-{client_addr[1]}"
                        )
                        self.receive_thread.start()
                        # Small delay to let thread start
                        time.sleep(0.01)
                        logger.info("✓ Receive thread started (thread ID: %s, alive: %s), waiting for handshake from %s...", self.receive_thread.ident, self.receive_thread.is_alive(), client_addr)
                    else:
                        logger.error("ERROR: client_socket is None after setting it - cannot start receive thread!")
                
            except socket.timeout:
                continue
            except OSError as e:
                if self.running:
                    logger.error("Error accepting connection: %s", e, exc_info=True)
                break
            except Exception as e:
                if self.running:
                    logger.error("Unexpected error in accept loop: %s", e, exc_info=True)
                break
    
    def _receive_messages_with_socket(self, socket_ref: socket.socket):
        """Receive messages from client using captured socket reference (runs in separate thread)"""
        thread_name = threading.current_thread().name
        logger.info("Started receiving messages from Electron (thread: %s)", thread_name)
        handshake_timer_start = time.time()
        
        while self.running:
            try:
                # Use captured socket reference to avoid race conditions
                logger.debug("Waiting to receive data from socket (thread: %s)...", thread_name)
                try:
                    data = socket_ref.recv(TCP_BUFFER_SIZE)
                except socket.timeout:
                    # Timeout is expected - check handshake timeout and continue
                    if not self.client_handshake_complete:
                        elapsed = time.time() - handshake_timer_start
                        if elapsed > HANDSHAKE_TIMEOUT:
                            logger.error("Handshake timeout after %ss - disconnecting client (thread: %s)", elapsed, thread_name)
                            break
                    logger.debug("Socket recv timeout (normal) - continuing (thread: %s)", thread_name)
                    continue
                
                if not data:
                    # Connection closed by client
                    logger.info("Electron client disconnected (no data received) (thread: %s)", thread_name)
                    break
                
                logger.info("✓ Received %s bytes from Electron (thread: %s)", len(data), thread_name)
                
                # Add to buffer
                self.receive_buffer += data
                
                # Process complete messages (newline-delimited)
                while b'\n' in self.receive_buffer:
                    line, self.receive_buffer = self.receive_buffer.split(b'\n', 1)
                    
                    if not line.strip():
                        continue
                    
                    try:
                        # Decode message
                        message_str = line.decode('utf-8').strip()
                        
                        # Check if this is a plain text idle popup response (format: "keep:1 continue" or "keep:0 stop")
                        if message_str.startswith('keep:'):
                            # Parse plain text format: "keep:1 continue" or "keep:0 stop"
                            try:
                                parts = message_str.split(' ', 1)
                                keep_part = parts[0]  # "keep:1" or "keep:0"
                                action = parts[1] if len(parts) > 1 else 'continue'  # "continue", "stop", etc.
                                
                                keep_value = keep_part.split(':')[1]  # "1" or "0"
                                keep_idle_time = (keep_value == '1')
                                
                                # Convert to JSON message format
                                message = {
                                    'type': 'idle_popup_response',
                                    'keep_idle_time': keep_idle_time,
                                    'action': action
                                }
                                message_type = 'idle_popup_response'
                                logger.info("✓ Received idle popup response from Electron: keep_idle_time=%s, action=%s", keep_idle_time, action)
                            except (ValueError, IndexError) as e:
                                logger.warning("Failed to parse idle popup response format: %s - %s", message_str, e)
                                continue
                        else:
                            # Try to parse as JSON
                            message = json.loads(message_str)
                            message_type = message.get('type', 'unknown')
                            logger.info("✓ Received message from Electron: type=%s", message_type)
                        
                        # Handle handshake
                        if message_type == 'hello_from_electron':
                            from ws_package.utils.ipc_session import (
                                verify_electron_handshake,
                                make_python_handshake_response,
                            )
                            accepted, secure = verify_electron_handshake(message)
                            if not accepted:
                                logger.error("IPC handshake auth failed — disconnecting client")
                                break
                            logger.info("✓ Handshake received: hello_from_electron")
                            # Send handshake response
                            with self.client_lock:
                                if self.client_socket is socket_ref:
                                    try:
                                        handshake_response = json.dumps(
                                            make_python_handshake_response(self.session_id, secure)
                                        ) + '\n'
                                        socket_ref.sendall(handshake_response.encode('utf-8'))
                                        self.client_handshake_complete = True
                                        logger.info("✓ Handshake complete - sent hello_from_python")
                                    except Exception as e:
                                        logger.error("Error sending handshake response: %s", e, exc_info=True)
                                        break
                            
                            # Send ready signal after handshake (outside the lock to avoid deadlock)
                            # Small delay to ensure Electron processes handshake
                            time.sleep(0.1)
                            try:
                                from ws_package.ws_engine import PYTHON_APP_VERSION
                                # Send ready message directly using the socket (handshake is complete)
                                ready_message = json.dumps({
                                    "type": "ready",
                                    "data": {
                                        "status": "ready",
                                        "message": "WS_Engine ready (TCP connected)",
                                        "capabilities": ["set_project_task", "stop"],
                                        "version": PYTHON_APP_VERSION
                                    }
                                }) + '\n'
                                socket_ref.sendall(ready_message.encode('utf-8'))
                                logger.info("✓ Sent 'ready' signal to Electron after handshake")
                            except Exception as e:
                                logger.error("Failed to send ready signal after handshake: %s", e, exc_info=True)
                            
                            # Don't queue handshake message, just process it
                            continue
                        
                        # Check handshake timeout
                        if not self.client_handshake_complete:
                            elapsed = time.time() - handshake_timer_start
                            if elapsed > HANDSHAKE_TIMEOUT:
                                logger.error("Handshake timeout after %ss - disconnecting client", elapsed)
                                break
                        
                        # Queue regular messages
                        try:
                            self.message_queue.put_nowait(message)
                            logger.debug("Queued message: type=%s, queue_size=%s", message_type, self.message_queue.qsize())
                        except queue.Full:
                            logger.warning("Message queue full (size=%s), dropping message: type=%s", self.message_queue.maxsize, message_type)
                        
                    except json.JSONDecodeError as e:
                        logger.warning("Invalid JSON from Electron: %s - %s", message_str[:200], e)
                    except UnicodeDecodeError as e:
                        logger.warning("Invalid encoding from Electron: %s", e)
                
            # Note: socket.timeout is already handled in the inner try/except block above
            # This outer except is for other exceptions
            except ConnectionResetError:
                logger.info("Connection reset by Electron client")
                break
            except BrokenPipeError:
                logger.info("Broken pipe - Electron client disconnected")
                break
            except OSError as e:
                if self.running:
                    logger.error("Socket error receiving message: %s", e, exc_info=True)
                break
            except Exception as e:
                if self.running:
                    logger.error("Error receiving message: %s", e, exc_info=True)
                break
        
        # Clean up client connection only if this socket is still the active one
        should_disconnect = False
        with self.client_lock:
            if self.client_socket is socket_ref:
                should_disconnect = True
        
        if should_disconnect:
            self._disconnect_client()
        logger.warning("Stopped receiving messages from Electron")
    
    def _receive_messages(self):
        """Legacy receive messages method - kept for backward compatibility"""
        with self.client_lock:
            socket_ref = self.client_socket
        if socket_ref:
            self._receive_messages_with_socket(socket_ref)
    
    def _disconnect_client(self):
        """Disconnect current client"""
        with self.client_lock:
            if self.client_socket:
                old_addr = self.client_address
                try:
                    self.client_socket.close()
                    logger.info("Closed client socket from %s", old_addr)
                except Exception as e:
                    logger.warning("Error closing client socket: %s", e)
                self.client_socket = None
                self.client_address = None
                self.client_handshake_complete = False
                self.receive_buffer = b''
                logger.info("Electron client disconnected from %s, waiting for reconnection...", old_addr)
    
    def send_message(self, message_type: str, data: Optional[Dict[str, Any]] = None) -> bool:
        """
        Send message to Electron client
        
        Args:
            message_type: Type of message (e.g., 'idle_popup_trigger', 'activity_update')
            data: Optional data dictionary to send
            
        Returns:
            True if sent successfully, False otherwise
        """
        # Capture socket reference to avoid race conditions
        with self.client_lock:
            socket_ref = self.client_socket
            handshake_complete = self.client_handshake_complete
        
        if not socket_ref:
            logger.debug("Cannot send message - no client connected: %s", message_type)
            return False
        
        if not handshake_complete:
            logger.debug("Cannot send message - handshake not complete: %s", message_type)
            return False
        
        # Pre-check: Verify socket is still connected using getpeername()
        try:
            socket_ref.getpeername()
        except (OSError, AttributeError):
            logger.debug("Cannot send message - socket not connected (getpeername failed): %s", message_type)
            return False
        
        message = {
            "type": message_type,
            "data": data
        }
        
        try:
            with self.send_lock:
                # Serialize to JSON and add newline
                message_str = json.dumps(message) + '\n'
                message_bytes = message_str.encode('utf-8')
                
                # Send message using captured socket reference
                socket_ref.sendall(message_bytes)
                logger.info("✓ Sent message to Electron: %s", message_type)
                return True
                
        except BrokenPipeError:
            logger.warning("Broken pipe while sending - client disconnected")
            self._disconnect_client()
            return False
        except ConnectionResetError:
            logger.warning("Connection reset while sending - client disconnected")
            self._disconnect_client()
            return False
        except OSError as e:
            logger.error("Socket error sending message: %s", e, exc_info=True)
            self._disconnect_client()
            return False
        except Exception as e:
            logger.error("Error sending message to Electron: %s", e, exc_info=True)
            return False
    
    def get_message(self, timeout: Optional[float] = None) -> Optional[Dict[str, Any]]:
        """
        Get a message from Electron (non-blocking or with timeout)
        
        Args:
            timeout: Timeout in seconds. If None, returns immediately if no message.
                     If 0, blocks indefinitely until message arrives.
        
        Returns:
            Message dictionary or None if timeout/empty
        """
        try:
            if timeout is None:
                return self.message_queue.get_nowait()
            else:
                return self.message_queue.get(timeout=timeout)
        except queue.Empty:
            return None
    
    def is_connected(self) -> bool:
        """
        Check if client is actually connected and handshake complete.
        Verifies the socket is still valid and connected.
        Non-destructive check - does not disconnect on failure.
        """
        with self.client_lock:
            socket_ref = self.client_socket
            handshake_complete = self.client_handshake_complete
        
        if not socket_ref or not handshake_complete:
            return False
        
        # Verify socket is actually connected by checking peer name
        try:
            socket_ref.getpeername()
            return True
        except (OSError, AttributeError):
            logger.debug("Socket exists but is not connected (getpeername failed)")
            return False
    
    def stop(self):
        """Stop the TCP server"""
        if not self.running:
            return
        
        logger.info("Stopping TCP server...")
        self.running = False
        
        # Disconnect client
        self._disconnect_client()
        
        # Close server socket with proper cleanup
        if self.server_socket:
            try:
                self.server_socket.shutdown(socket.SHUT_RDWR)
            except Exception:
                pass
            try:
                self.server_socket.close()
            except Exception:
                pass
            self.server_socket = None
        
        time.sleep(0.1)
        
        logger.info("TCP server stopped")


# Global instance
_tcp_server: Optional[TCPServer] = None


def get_tcp_server() -> TCPServer:
    """Get the global TCP server instance"""
    global _tcp_server
    if _tcp_server is None:
        _tcp_server = TCPServer()
    return _tcp_server


def start_tcp_server():
    """Start the TCP server (call this at startup)"""
    server = get_tcp_server()
    server.start()


def stop_tcp_server():
    """Stop the TCP server (call this at shutdown)"""
    global _tcp_server
    if _tcp_server:
        _tcp_server.stop()
        _tcp_server = None
