"""
Engine Module - Unified management for Register and Interceptor

Provides a simple interface to start Register and/or Interceptor
for user services.

Usage:
```python
from gmi_ieops.handler import Engine
from gmi_ieops.utils import get_service_port

port = get_service_port()

# Start engine (auto cleanup on exit)
# Automatically waits for model sync if MODEL_FROM is set
engine = Engine(
    service_port=port,
    register_enable=True,
    interceptor_enable=True,
    wait_model_sync=True,  # Auto wait for model sync (default: True)
)
engine.start()

# Start your service
app.run(host="127.0.0.1", port=port)
```
"""

import os
import sys
import atexit
import signal
import asyncio
import multiprocessing
from typing import Optional, Callable, Any, Set

from ..utils.log import uvicorn_logger
from ..utils.util import get_socket_path
from ..config import env


# Global engine process for cleanup
_engine_process: Optional[multiprocessing.Process] = None
_original_sigint_handler: Optional[Callable] = None
_original_sigterm_handler: Optional[Callable] = None


def _cleanup_engine():
    """Cleanup engine subprocess on exit"""
    global _engine_process
    if _engine_process is not None and _engine_process.is_alive():
        uvicorn_logger.info("Stopping engine subprocess...")
        _engine_process.terminate()
        _engine_process.join(timeout=5)
        if _engine_process.is_alive():
            _engine_process.kill()
            _engine_process.join(timeout=1)
        _engine_process = None


def _signal_handler(signum, frame):
    """Handle signals - cleanup engine then call original handler"""
    global _original_sigint_handler, _original_sigterm_handler
    
    # Cleanup engine subprocess
    _cleanup_engine()
    
    # Call original handler if exists
    if signum == signal.SIGINT and _original_sigint_handler:
        if callable(_original_sigint_handler):
            _original_sigint_handler(signum, frame)
        elif _original_sigint_handler == signal.SIG_DFL:
            signal.signal(signum, signal.SIG_DFL)
            os.kill(os.getpid(), signum)
    elif signum == signal.SIGTERM and _original_sigterm_handler:
        if callable(_original_sigterm_handler):
            _original_sigterm_handler(signum, frame)
        elif _original_sigterm_handler == signal.SIG_DFL:
            signal.signal(signum, signal.SIG_DFL)
            os.kill(os.getpid(), signum)
    else:
        # No original handler, use default behavior
        signal.signal(signum, signal.SIG_DFL)
        os.kill(os.getpid(), signum)


# Register cleanup handlers
atexit.register(_cleanup_engine)


class Engine:
    """
    Unified engine for Register and Interceptor management.
    
    Features:
    - Optionally enable Register (service registration with IEOPS proxy)
    - Optionally enable Interceptor (transparent request proxy)
    - Automatic cleanup on exit
    - Runs in subprocess, user service runs in main process
    - Auto wait for model sync before starting (if MODEL_FROM is set)
    
    Example:
        from gmi_ieops.handler import Engine
        from gmi_ieops.utils import get_service_port
        
        port = get_service_port()
        engine = Engine(service_port=port)
        engine.start()  # Auto waits for model sync if MODEL_FROM is set
        app.run(host="127.0.0.1", port=port)
    """
    
    def __init__(
        self,
        service_port: Optional[int] = None,
        service_host: str = "127.0.0.1",
        register_enable: bool = True,
        interceptor_enable: bool = True,
        timeout: int = 600,
        app_name: Optional[str] = None,
        socket_path: Optional[str] = None,
        filter_headers: Optional[Set[str]] = None,
        filter_response_headers: Optional[Set[str]] = None,
        wait_model_sync: bool = True,
        model_sync_timeout: Optional[int] = None,
        model_sync_interval: Optional[int] = None,
        exit_on_sync_fail: bool = True,
    ) -> None:
        """
        Initialize engine.
        
        Args:
            service_port: Port where user service is listening.
                          Required if interceptor_enable=True.
            service_host: Host where user service is listening (default: 127.0.0.1)
            register_enable: Enable service registration with IEOPS proxy
            interceptor_enable: Enable transparent request proxy
            timeout: Request timeout in seconds (for interceptor)
            app_name: Application name (defaults to env.app.name)
            socket_path: Unix socket path (defaults to auto-generated)
            filter_headers: Set of header names (lowercase) to filter from requests.
                           Defaults to DEFAULT_FILTER_HEADERS.
            filter_response_headers: Set of header names (lowercase) to filter from responses.
                           Defaults to DEFAULT_FILTER_RESPONSE_HEADERS.
            wait_model_sync: If True, wait for model sync completion before starting.
                            Only applies when MODEL_FROM env var is set. (default: True)
            model_sync_timeout: Max wait time for model sync in seconds.
                               Defaults to MODEL_SYNC_WAIT_TIMEOUT env var (24h).
            model_sync_interval: Check interval for model sync in seconds.
                                Defaults to MODEL_SYNC_WAIT_INTERVAL env var (5s).
            exit_on_sync_fail: If True, exit process when model sync fails. (default: True)
        """
        self._service_port = service_port
        self._service_host = service_host
        self._register_enable = register_enable
        self._interceptor_enable = interceptor_enable
        self._timeout = timeout
        self._app_name = app_name or env.app.name
        self._socket_path = socket_path or get_socket_path()
        self._filter_headers = filter_headers
        self._filter_response_headers = filter_response_headers
        self._wait_model_sync = wait_model_sync
        self._model_sync_timeout = model_sync_timeout
        self._model_sync_interval = model_sync_interval
        self._exit_on_sync_fail = exit_on_sync_fail
        
        if interceptor_enable and not service_port:
            raise ValueError("service_port is required when interceptor_enable=True")
        
        # Set metrics port if not explicitly configured
        if service_port and not env.metrics.port:
            env.metrics.port = service_port
    
    def start(self) -> bool:
        """
        Start engine in subprocess.
        
        The subprocess runs Register and/or Interceptor based on configuration.
        Automatically cleaned up on program exit or signals.
        
        Before starting, if wait_model_sync=True and MODEL_FROM is set,
        this method will block until the model sync completes.
        
        Returns:
            True if started successfully, False if model sync failed (when exit_on_sync_fail=False).
        """
        global _engine_process, _original_sigint_handler, _original_sigterm_handler
        
        # Wait for model sync if enabled and MODEL_FROM is set
        if self._wait_model_sync and env.model_sync.model_from:
            if not self._wait_for_model_sync():
                if self._exit_on_sync_fail:
                    uvicorn_logger.error("Model sync failed! Exiting process...")
                    sys.exit(1)
                return False
        
        _cleanup_engine()
        
        if not self._register_enable and not self._interceptor_enable:
            uvicorn_logger.info("Engine: both register and interceptor disabled, nothing to start")
            return True
        
        proc = multiprocessing.Process(
            target=_run_engine,
            args=(
                self._service_port,
                self._service_host,
                self._register_enable,
                self._interceptor_enable,
                self._timeout,
                self._app_name,
                self._socket_path,
                self._filter_headers,
                self._filter_response_headers,
            ),
            daemon=False,  # Non-daemon so it can cleanup properly
        )
        proc.start()
        
        _engine_process = proc
        
        # Save original handlers and install ours
        try:
            _original_sigint_handler = signal.signal(signal.SIGINT, _signal_handler)
            _original_sigterm_handler = signal.signal(signal.SIGTERM, _signal_handler)
        except ValueError:
            # Not in main thread
            pass
        
        modes = []
        if self._register_enable:
            modes.append("register")
        if self._interceptor_enable:
            modes.append("interceptor")
        
        uvicorn_logger.info(f"Engine started (PID: {proc.pid}), modes: {', '.join(modes)}")
        if self._interceptor_enable:
            uvicorn_logger.info(f"Forwarding to: http://{self._service_host}:{self._service_port}")
        
        return True
    
    def _wait_for_model_sync(self) -> bool:
        """
        Wait for model sync to complete.
        
        Returns:
            True if sync completed successfully, False otherwise.
        """
        from .model_sync import wait_for_model_sync
        
        model_from = env.model_sync.model_from
        model_path = env.model_sync.model_path
        
        uvicorn_logger.info(f"Model sync enabled - MODEL_FROM: {model_from}")
        uvicorn_logger.info(f"Model sync enabled - MODEL_PATH: {model_path}")
        uvicorn_logger.info("Waiting for model sync to complete before starting engine...")
        
        success = wait_for_model_sync(
            timeout=self._model_sync_timeout,
            interval=self._model_sync_interval,
            logger=uvicorn_logger,
        )
        
        if success:
            uvicorn_logger.info("Model sync completed successfully!")
        else:
            uvicorn_logger.error("Model sync failed or timed out!")
        
        return success


def _run_engine(
    service_port: Optional[int],
    service_host: str,
    register_enable: bool,
    interceptor_enable: bool,
    timeout: int,
    app_name: str,
    socket_path: str,
    filter_headers: Optional[Set[str]],
    filter_response_headers: Optional[Set[str]],
) -> None:
    """Run engine in subprocess"""
    # Setup signal handlers for graceful shutdown
    shutdown_event = asyncio.Event()
    
    def handle_signal(signum, frame):
        uvicorn_logger.info(f"Engine subprocess received signal {signum}")
        shutdown_event.set()
    
    signal.signal(signal.SIGTERM, handle_signal)
    signal.signal(signal.SIGINT, handle_signal)
    
    asyncio.run(_async_run_engine(
        service_port,
        service_host,
        register_enable,
        interceptor_enable,
        timeout,
        app_name,
        socket_path,
        filter_headers,
        filter_response_headers,
        shutdown_event,
    ))


async def _async_run_engine(
    service_port: Optional[int],
    service_host: str,
    register_enable: bool,
    interceptor_enable: bool,
    timeout: int,
    app_name: str,
    socket_path: str,
    filter_headers: Optional[Set[str]],
    filter_response_headers: Optional[Set[str]],
    shutdown_event: asyncio.Event,
) -> None:
    """Async engine runner - manages Register and Interceptor"""
    from .interceptor import Interceptor
    from .register import Register
    
    uvicorn_logger.info(f"{app_name} Engine starting...")
    
    tasks = []
    interceptor_runner = None
    
    # Start Register if enabled
    if register_enable:
        register = Register()
        register_task = asyncio.create_task(register.register())
        tasks.append(register_task)
        uvicorn_logger.info("Register enabled")
    
    # Start Interceptor if enabled
    if interceptor_enable:
        upstream = f"http://{service_host}:{service_port}"
        interceptor = Interceptor(
            upstream=upstream,
            timeout=timeout,
            app_name=app_name,
            socket_path=socket_path,
            filter_headers=filter_headers,
            filter_response_headers=filter_response_headers,
        )
        interceptor_runner = await interceptor.serve_async()
        uvicorn_logger.info(f"Interceptor enabled, upstream: {upstream}")
    
    # Wait for shutdown signal
    await shutdown_event.wait()
    
    uvicorn_logger.info(f"{app_name} Engine shutting down...")
    
    # Cleanup Register tasks
    for task in tasks:
        if not task.done():
            task.cancel()
            try:
                await task
            except asyncio.CancelledError:
                pass
    
    # Cleanup Interceptor
    if interceptor_runner:
        await interceptor_runner.cleanup()
    
    uvicorn_logger.info(f"{app_name} Engine stopped")
