# Generated by Copilot
# Maintained by Dinesh Kumar C (https://github.com/Dinesh210805/blender-mcp)
# Forked from original work by Siddharth Ahuja (https://github.com/ahujasid/blender-mcp)

"""Enhanced connection management for Blender MCP."""

import socket
import json
import time
import threading
from typing import Dict, Any, Optional
from dataclasses import dataclass, field

from .config import config
from .exceptions import ConnectionError, TimeoutError, BlenderOperationError
from .utils import get_logger

logger = get_logger("Connection")


@dataclass
class BlenderConnection:
    """Enhanced Blender connection with proper error handling and retry logic."""
    
    host: str = field(default_factory=lambda: config.blender_host)
    port: int = field(default_factory=lambda: config.blender_port)
    timeout: float = field(default_factory=lambda: config.socket_timeout)
    retry_attempts: int = field(default_factory=lambda: config.connection_retry_attempts)
    retry_delay: float = field(default_factory=lambda: config.connection_retry_delay)
    
    _sock: Optional[socket.socket] = field(default=None, init=False)
    _lock: threading.RLock = field(default_factory=threading.RLock, init=False)
    _connection_count: int = field(default=0, init=False)
    
    def __post_init__(self):
        """Initialize after dataclass creation."""
        self._connection_id = id(self)
        logger.debug(f"Created BlenderConnection {self._connection_id}")
    
    def connect(self) -> bool:
        """Connect to Blender with retry logic."""
        with self._lock:
            if self._sock and self._is_socket_alive():
                logger.debug("Using existing connection")
                return True
            
            self._close_socket()
            
            for attempt in range(self.retry_attempts):
                try:
                    logger.info(f"Connecting to Blender at {self.host}:{self.port} (attempt {attempt + 1}/{self.retry_attempts})")
                    
                    self._sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
                    self._sock.settimeout(self.timeout)
                    self._sock.connect((self.host, self.port))
                    
                    self._connection_count += 1
                    logger.info(f"Successfully connected to Blender (connection #{self._connection_count})")
                    return True
                    
                except (socket.error, OSError) as e:
                    logger.warning(f"Connection attempt {attempt + 1} failed: {e}")
                    self._close_socket()
                    
                    if attempt < self.retry_attempts - 1:
                        logger.info(f"Retrying in {self.retry_delay} seconds...")
                        time.sleep(self.retry_delay)
            
            logger.error(f"Failed to connect to Blender after {self.retry_attempts} attempts")
            raise ConnectionError(f"Could not connect to Blender at {self.host}:{self.port} after {self.retry_attempts} attempts")
    
    def disconnect(self):
        """Disconnect from Blender."""
        with self._lock:
            self._close_socket()
            logger.info("Disconnected from Blender")
    
    def _close_socket(self):
        """Internal method to close socket."""
        if self._sock:
            try:
                self._sock.close()
            except Exception as e:
                logger.debug(f"Error closing socket: {e}")
            finally:
                self._sock = None
    
    def _is_socket_alive(self) -> bool:
        """Check if socket connection is alive."""
        if not self._sock:
            return False
        
        try:
            # Use SO_ERROR to check socket status
            error = self._sock.getsockopt(socket.SOL_SOCKET, socket.SO_ERROR)
            return error == 0
        except Exception:
            return False
    
    def _receive_complete_response(self) -> bytes:
        """Receive complete JSON response from socket."""
        if not self._sock:
            raise ConnectionError("Not connected to Blender")
        
        chunks = []
        start_time = time.time()
        
        try:
            while True:
                if time.time() - start_time > self.timeout:
                    raise TimeoutError(f"Timeout ({self.timeout}s) while receiving response")
                
                try:
                    chunk = self._sock.recv(config.socket_buffer_size)
                    if not chunk:
                        if not chunks:
                            raise ConnectionError("Connection closed before receiving any data")
                        break
                    
                    chunks.append(chunk)
                    
                    # Try to parse accumulated data as JSON
                    try:
                        data = b''.join(chunks)
                        json.loads(data.decode('utf-8'))
                        logger.debug(f"Received complete response ({len(data)} bytes)")
                        return data
                    except json.JSONDecodeError:
                        # Incomplete JSON, continue receiving
                        continue
                        
                except socket.timeout:
                    logger.warning("Socket timeout during receive")
                    break
                except (ConnectionError, BrokenPipeError, ConnectionResetError) as e:
                    logger.error(f"Socket connection error during receive: {e}")
                    raise ConnectionError(f"Connection lost during receive: {e}")
        
        except Exception as e:
            logger.error(f"Error during receive: {e}")
            self._close_socket()
            raise
        
        # Try to use what we have
        if chunks:
            data = b''.join(chunks)
            try:
                json.loads(data.decode('utf-8'))
                return data
            except json.JSONDecodeError:
                raise BlenderOperationError("Incomplete JSON response received")
        
        raise BlenderOperationError("No data received from Blender")
    
    def send_command(self, command_type: str, params: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
        """Send command to Blender with proper error handling."""
        with self._lock:
            if not self._sock and not self.connect():
                raise ConnectionError("Cannot establish connection to Blender")
            
            command = {
                "type": command_type,
                "params": params or {}
            }
            
            start_time = time.time()
            
            try:
                logger.debug(f"Sending command: {command_type}")
                
                # Send command
                command_json = json.dumps(command).encode('utf-8')
                self._sock.sendall(command_json)
                
                # Receive response
                response_data = self._receive_complete_response()
                response = json.loads(response_data.decode('utf-8'))
                
                duration = time.time() - start_time
                logger.debug(f"Command {command_type} completed in {duration:.2f}s")
                
                # Check for errors in response
                if response.get("status") == "error":
                    error_msg = response.get('message', 'Unknown error from Blender')
                    logger.error(f"Blender error for {command_type}: {error_msg}")
                    raise BlenderOperationError(error_msg, error_code="BLENDER_ERROR")
                
                return response.get("result", {})
                
            except socket.timeout:
                logger.error(f"Timeout waiting for response to {command_type}")
                self._close_socket()
                raise TimeoutError(f"Timeout waiting for Blender response to {command_type}")
                
            except (ConnectionError, BrokenPipeError, ConnectionResetError) as e:
                logger.error(f"Connection error during {command_type}: {e}")
                self._close_socket()
                raise ConnectionError(f"Connection lost during {command_type}: {e}")
                
            except json.JSONDecodeError as e:
                logger.error(f"Invalid JSON response for {command_type}: {e}")
                raise BlenderOperationError(f"Invalid response from Blender for {command_type}: {e}")
                
            except Exception as e:
                logger.error(f"Unexpected error during {command_type}: {e}")
                self._close_socket()
                raise BlenderOperationError(f"Unexpected error during {command_type}: {e}")
    
    def __enter__(self):
        """Context manager entry."""
        self.connect()
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        """Context manager exit."""
        self.disconnect()


class ConnectionManager:
    """Manages Blender connections with proper lifecycle."""
    
    def __init__(self):
        self._connection: Optional[BlenderConnection] = None
        self._lock = threading.RLock()
    
    def get_connection(self) -> BlenderConnection:
        """Get or create a Blender connection."""
        with self._lock:
            if self._connection is None:
                logger.info("Creating new Blender connection")
                self._connection = BlenderConnection()
            
            try:
                self._connection.connect()
                return self._connection
            except Exception as e:
                logger.error(f"Failed to establish connection: {e}")
                self._connection = None
                raise
    
    def close_connection(self):
        """Close the current connection."""
        with self._lock:
            if self._connection:
                self._connection.disconnect()
                self._connection = None
                logger.info("Connection closed")
    
    def execute_command(self, command_type: str, params: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
        """Execute a command using the managed connection."""
        connection = self.get_connection()
        return connection.send_command(command_type, params)


# Global connection manager instance
_connection_manager = ConnectionManager()


def get_blender_connection() -> BlenderConnection:
    """Get the global Blender connection."""
    return _connection_manager.get_connection()


def execute_blender_command(command_type: str, params: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
    """Execute a Blender command using the global connection manager."""
    return _connection_manager.execute_command(command_type, params)


def close_blender_connection():
    """Close the global Blender connection."""
    _connection_manager.close_connection()
