"""
HTTP server for hosting shared functions
"""

import json
import pickle
import threading
from http.server import HTTPServer, ThreadingHTTPServer, BaseHTTPRequestHandler
from typing import Callable, Dict
import traceback

# Global registry of shared functions
_registry: Dict[str, Callable] = {}
_server = None
_server_thread = None
_server_port = 5555
_lock = threading.Lock()


def register_function(name: str, func: Callable):
    """Register a function to be callable remotely"""
    _registry[name] = func


def get_registry():
    """Get the function registry"""
    return _registry


class RPCHandler(BaseHTTPRequestHandler):
    """HTTP request handler for RPC calls"""
    
    def log_message(self, format, *args):
        """Suppress default logging"""
        pass
    
    def do_POST(self):
        """Handle POST requests for function calls"""
        try:
            content_length = int(self.headers['Content-Length'])
            body = self.rfile.read(content_length)
            
            # Parse request
            request = json.loads(body.decode('utf-8'))
            func_name = request['function']
            args = request.get('args', [])
            kwargs = request.get('kwargs', {})
            
            # Deserialize pickled arguments if present
            if 'args_pickle' in request:
                args = pickle.loads(bytes.fromhex(request['args_pickle']))
            if 'kwargs_pickle' in request:
                kwargs = pickle.loads(bytes.fromhex(request['kwargs_pickle']))
            
            # Get function from registry
            if func_name not in _registry:
                self.send_error(404, f"Function '{func_name}' not found")
                return
            
            func = _registry[func_name]
            
            # Handle the case where the server script is run as __main__
            # We want pickled objects to have the correct module name (e.g., 'tests.daemon_for_testing')
            # instead of '__main__', so they can be unpickled by other processes.
            import sys
            import inspect
            func_module = inspect.getmodule(func)
            original_module_names = {} # To restore __module__ on classes
            
            if func_module and func_module.__name__ == '__main__':
                # Extract the real module name from the registered function name
                parts = func_name.rsplit('.', 1)
                if len(parts) > 1:
                    real_module_name = parts[0]
                    
                    # Temporarily update __module__ of all classes in this module
                    # so pickle uses the real name instead of '__main__'
                    for name, obj in inspect.getmembers(func_module):
                        if inspect.isclass(obj) and obj.__module__ == '__main__':
                            original_module_names[obj] = obj.__module__
                            obj.__module__ = real_module_name
                    
                    # Also register it in sys.modules under its real name if not already there
                    if real_module_name not in sys.modules:
                        sys.modules[real_module_name] = func_module

            try:
                # Call function
                result = func(*args, **kwargs)
                
                # Serialize result
                try:
                    # Try JSON first (faster and more portable)
                    json.dumps(result)
                    response = {
                        'success': True,
                        'result': result,
                        'serialization': 'json'
                    }
                    response_body = json.dumps(response).encode('utf-8')
                except (TypeError, ValueError):
                    # Fall back to pickle for complex objects
                    response = {
                        'success': True,
                        'result_pickle': pickle.dumps(result).hex(),
                        'serialization': 'pickle'
                    }
                    response_body = json.dumps(response).encode('utf-8')
            finally:
                # Restore original __module__ on classes
                for obj, original_name in original_module_names.items():
                    obj.__module__ = original_name

            # Send response
            self.send_response(200)
            self.send_header('Content-Type', 'application/json')
            self.send_header('Content-Length', str(len(response_body)))
            self.end_headers()
            self.wfile.write(response_body)
            
        except Exception as e:
            error_msg = f"{type(e).__name__}: {str(e)}"
            error_trace = traceback.format_exc()
            
            response = {
                'success': False,
                'error': error_msg,
                'traceback': error_trace
            }
            response_body = json.dumps(response).encode('utf-8')
            
            self.send_response(500)
            self.send_header('Content-Type', 'application/json')
            self.send_header('Content-Length', str(len(response_body)))
            self.end_headers()
            self.wfile.write(response_body)
    
    def do_GET(self):
        """Handle GET requests for server info"""
        if self.path == '/health':
            response = {'status': 'ok', 'functions': list(_registry.keys())}
            response_body = json.dumps(response).encode('utf-8')
            
            self.send_response(200)
            self.send_header('Content-Type', 'application/json')
            self.send_header('Content-Length', str(len(response_body)))
            self.end_headers()
            self.wfile.write(response_body)
        else:
            self.send_error(404)


def start_server(port: int = 5555):
    """Start the HTTP server in a background thread"""
    global _server, _server_thread, _server_port
    
    with _lock:
        if _server is not None:
            return  # Server already running
        
        _server_port = port
        _server = ThreadingHTTPServer(('localhost', port), RPCHandler)
        
        def serve():
            print(f"[remotefunc] Server started on http://localhost:{port}")
            _server.serve_forever()
        
        _server_thread = threading.Thread(target=serve, daemon=True)
        _server_thread.start()


def shutdown_server():
    """Shutdown the HTTP server"""
    global _server, _server_thread
    
    with _lock:
        if _server is not None:
            _server.shutdown()
            _server = None
            _server_thread = None
            print("[remotefunc] Server shutdown")


def get_server_port():
    """Get the current server port"""
    return _server_port


def is_server_running():
    """Check if server is running"""
    return _server is not None
