content_server.py

import socket
import threading
import os
import time
from typing import Dict, List, Optional, Tuple, Any
from pathlib import Path


class ContentServer:
    
    def __init__(
        self,
        server_id: str,
        ip_address: str = 'localhost',
        index_server_host: str = 'localhost',
        index_server_port: int = 5000,
        tcp_port: int = 7001,
        udp_port: int = 7002,
        monitor_host: str = 'localhost',
        monitor_udp_port: int = 6000,
        heartbeat_interval: int = 3,
        files_directory: str = None
    ) -> None:
        self.server_id = server_id
        
        logs_dir = Path("logs")
        logs_dir.mkdir(exist_ok=True)
        self.log_file = str(logs_dir / f"{server_id}_log.txt")
        self.log_lock: threading.Lock = threading.Lock()
        
        with open(self.log_file, 'w', encoding='utf-8') as f:
            f.write("")
        
        if ip_address == 'localhost' or ip_address is None:
            self.ip_address = self._get_local_ip()
        else:
            self.ip_address = ip_address
        self.index_server_host = index_server_host
        self.index_server_port = index_server_port
        self.tcp_port = tcp_port
        self.udp_port = udp_port
        self.monitor_host = monitor_host
        self.monitor_udp_port = monitor_udp_port
        self.heartbeat_interval = heartbeat_interval
        
        self.files: Dict[str, Dict[str, Any]] = {}
        self.files_directory = files_directory
        
        if files_directory:
            self._load_from_directory(files_directory)
        else:
            self._log_message("ERROR", "files_directory parameter is required")
        
        self.tcp_socket: Optional[socket.socket] = None
        self.udp_socket: Optional[socket.socket] = None
        self.running: bool = False
        
        self.active_clients: int = 0
        self.client_lock: threading.Lock = threading.Lock()
        
        self.lock: threading.Lock = threading.Lock()
    
    def _log_message(self, direction: str, message: str, target: Optional[str] = None) -> None:
        try:
            with self.log_lock:
                timestamp = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())
                target_info = f" to {target}" if direction == "SEND" and target else (f" from {target}" if direction == "RECEIVE" and target else "")
                with open(self.log_file, 'a', encoding='utf-8') as f:
                    f.write(f"[{timestamp}] {direction}{target_info}: {message}\n")
        except Exception as e:
            print(f"Error writing to log file: {e}")
    
    def _get_local_ip(self) -> str:
        try:
            temp_sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
            temp_sock.connect(('8.8.8.8', 80))
            local_ip = temp_sock.getsockname()[0]
            temp_sock.close()
            self._log_message("INFO", f"Auto-detected IP address: {local_ip}")
            return local_ip
        except Exception as e:
            self._log_message("INFO", f"Warning: Could not auto-detect IP address: {e}. Using 'localhost'")
            return 'localhost'
    
    def _load_from_directory(self, directory: str) -> None:
        dir_path = Path(directory)
        if not dir_path.exists():
            self._log_message("ERROR", f"Directory not found: {directory}")
            return
        
        self._log_message("INFO", f"Loading files from directory: {directory}")
        for file_path in dir_path.iterdir():
            if file_path.is_file():
                filename = file_path.name
                size = file_path.stat().st_size
                self.files[filename] = {
                    'path': str(file_path),
                    'size': size
                }
                self._log_message("INFO", f"Found file: {filename} ({size} bytes)")
    
    def _get_file_data(self, filename: str) -> Optional[bytes]:
        if filename not in self.files:
            return None
        
        file_info = self.files[filename]
        
        if file_info['path'] and os.path.exists(file_info['path']):
            with open(file_info['path'], 'rb') as f:
                return f.read()
        
        return None
    
    def register_with_index_server(self) -> bool:
        try:
            sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            sock.connect((self.index_server_host, self.index_server_port))
            
            register_msg = f"REGISTER {self.server_id} {self.tcp_port} {self.udp_port}\n"
            sock.send(register_msg.encode('utf-8'))
            self._log_message("SEND", register_msg.strip(), "Index Server")
            self._log_message("INFO", f"Sent REGISTER to Index Server: {register_msg.strip()}")
            
            response = sock.recv(1024).decode('utf-8').strip()
            self._log_message("RECEIVE", response, "Index Server")
            self._log_message("INFO", f"Received from Index Server: {response}")
            
            if response != "OK REGISTERED":
                self._log_message("ERROR", f"Registration failed: {response}")
                sock.close()
                return False
            
            for filename, file_info in self.files.items():
                add_file_msg = f"ADD_FILE {self.server_id} {filename} {file_info['size']}\n"
                sock.send(add_file_msg.encode('utf-8'))
                self._log_message("SEND", add_file_msg.strip(), "Index Server")
                self._log_message("INFO", f"Sent: {add_file_msg.strip()}")
            
            done_msg = "DONE_FILES\n"
            sock.send(done_msg.encode('utf-8'))
            self._log_message("SEND", done_msg.strip(), "Index Server")
            self._log_message("INFO", f"Sent: {done_msg.strip()}")
            
            response = sock.recv(1024).decode('utf-8').strip()
            self._log_message("RECEIVE", response, "Index Server")
            self._log_message("INFO", f"Received from Index Server: {response}")
            
            if response == "OK FILES_ADDED":
                self._log_message("INFO", "Successfully registered with Index Server")
                sock.close()
                return True
            else:
                self._log_message("ERROR", f"File list registration failed: {response}")
                sock.close()
                return False
                
        except Exception as e:
            self._log_message("ERROR", f"Error registering with Index Server: {e}")
            return False
    
    def _read_line(self, client_socket: socket.socket) -> Optional[str]:
        buffer = ""
        try:
            while True:
                data = client_socket.recv(1).decode('utf-8')
                if not data:
                    return None
                if data == '\n':
                    line = buffer.strip()
                    return line
                buffer += data
        except (ConnectionResetError, ConnectionAbortedError, BrokenPipeError):
            return None
        except socket.error as e:
            self._log_message("ERROR", f"Socket error while reading: {e}")
            return None
        except UnicodeDecodeError:
            self._log_message("ERROR", "Invalid encoding in message")
            return None
    
    def _handle_client_request(self, client_socket: socket.socket, address: Tuple[str, int]) -> None:
        try:
            with self.client_lock:
                self.active_clients += 1
            
            while True:
                line = self._read_line(client_socket)
                if not line:
                    self._log_message("INFO", f"Client {address} disconnected")
                    break
                
                self._log_message("RECEIVE", line, "Client")
                
                if not line.startswith('GET '):
                    try:
                        error_msg = "ERROR INVALID_COMMAND\n"
                        client_socket.send(error_msg.encode('utf-8'))
                        self._log_message("SEND", error_msg.strip(), "Client")
                        self._log_message("INFO", f"Invalid command from {address}: {line}")
                    except:
                        pass
                    break
                
                parts = line.split(None, 1)
                if len(parts) != 2:
                    try:
                        error_msg = "ERROR INVALID_COMMAND\n"
                        client_socket.send(error_msg.encode('utf-8'))
                        self._log_message("SEND", error_msg.strip(), "Client")
                        self._log_message("INFO", f"Malformed GET command from {address}: {line}")
                    except:
                        pass
                    break
                
                filename = parts[1]
                self._log_message("INFO", f"Client {address} requested file: {filename}")
                
                file_data = self._get_file_data(filename)
                
                if file_data is None:
                    try:
                        error_msg = "ERROR FILE_NOT_FOUND\n"
                        client_socket.send(error_msg.encode('utf-8'))
                        self._log_message("SEND", error_msg.strip(), "Client")
                        self._log_message("INFO", f"File not found: {filename}")
                    except:
                        pass
                    break
                else:
                    try:
                        ok_msg = f"OK {len(file_data)}\n"
                        client_socket.send(ok_msg.encode('utf-8'))
                        self._log_message("SEND", ok_msg.strip(), "Client")
                        self._log_message("INFO", f"File found: {filename} ({len(file_data)} bytes)")
                    except:
                        self._log_message("INFO", f"Client {address} disconnected before file transfer")
                        break
                    
                    chunk_size = 4096
                    total_sent = 0
                    try:
                        while total_sent < len(file_data):
                            chunk = file_data[total_sent:total_sent + chunk_size]
                            sent = client_socket.send(chunk)
                            if sent == 0:
                                self._log_message("INFO", "Connection closed by client during transfer")
                                break
                            total_sent += sent
                        self._log_message("INFO", f"Sent file {filename} ({total_sent} bytes) to {address}")
                    except Exception as e:
                        self._log_message("INFO", f"Client {address} disconnected during file transfer: {e}")
                        self._log_message("INFO", f"Sent {total_sent}/{len(file_data)} bytes before disconnection")
                        break
                
        except Exception as e:
            self._log_message("ERROR", f"Unexpected error handling client request from {address}: {e}")
        finally:
            try:
                with self.client_lock:
                    self.active_clients -= 1
                client_socket.close()
                self._log_message("INFO", f"Client connection closed: {address}")
            except:
                pass
    
    def _send_heartbeat(self) -> None:
        while self.running:
            try:
                time.sleep(self.heartbeat_interval)
                
                if not self.running:
                    break
                
                with self.client_lock:
                    current_load = self.active_clients
                
                num_files = len(self.files)
                
                heartbeat_msg = f"HEARTBEAT {self.server_id} {self.ip_address} {self.tcp_port} {current_load} {num_files}\n"
                
                if self.udp_socket:
                    self.udp_socket.sendto(
                        heartbeat_msg.encode('utf-8'),
                        (self.monitor_host, self.monitor_udp_port)
                    )
                    self._log_message("SEND", heartbeat_msg.strip(), "Monitor Server")
                
            except Exception as e:
                if self.running:
                    self._log_message("ERROR", f"Error sending heartbeat: {e}")
    
    def start(self) -> None:
        if not self.register_with_index_server():
            self._log_message("ERROR", "Failed to register with Index Server. Exiting.")
            return
        
        self.udp_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        try:
            self.udp_socket.bind(('0.0.0.0', self.udp_port))
        except OSError:
            self.udp_socket.bind(('0.0.0.0', 0))
            self.udp_port = self.udp_socket.getsockname()[1]
        
        self.running = True
        
        heartbeat_thread = threading.Thread(target=self._send_heartbeat, daemon=True)
        heartbeat_thread.start()
        self._log_message("INFO", f"Heartbeat thread started (interval: {self.heartbeat_interval}s)")
        
        self.tcp_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.tcp_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        self.tcp_socket.bind(('0.0.0.0', self.tcp_port))
        self.tcp_socket.listen(10)
        self._log_message("INFO", f"Content Server {self.server_id} started on TCP port {self.tcp_port} - Serving {len(self.files)} files, Monitor: {self.monitor_host}:{self.monitor_udp_port}")
        
        while self.running:
            try:
                client_socket, address = self.tcp_socket.accept()
                self._log_message("INFO", f"New client connection from {address} (Total active: {self.active_clients + 1})")
                
                thread = threading.Thread(
                    target=self._handle_client_request,
                    args=(client_socket, address),
                    daemon=True
                )
                thread.start()
                
            except Exception as e:
                if self.running:
                    self._log_message("ERROR", f"Error accepting connection: {e}")
    
    def stop(self) -> None:
        self.running = False
        if self.tcp_socket:
            self.tcp_socket.close()
        if self.udp_socket:
            self.udp_socket.close()
        self._log_message("INFO", f"Content Server {self.server_id} stopped")


if __name__ == '__main__':
    print("Content Server should be started via start_server1.py or start_server2.py")