# ruff: noqa: E402
import asyncio
import inspect
import os
import signal
import logging
import ipaddress
import threading
import time
import uuid
from pathlib import Path

from sari.mcp.session import Session
from sari.core.workspace import WorkspaceManager
from sari.core.server_registry import ServerRegistry
from sari.core.settings import settings
from sari.core.constants import DEFAULT_DAEMON_HOST, DEFAULT_DAEMON_PORT
from sari.mcp.trace import trace

def _resolve_log_dir() -> Path:
    val = settings.LOG_DIR
    if val:
        return Path(os.path.expanduser(val)).resolve()
    return WorkspaceManager.get_global_log_dir()

def _init_logging() -> None:
    log_dir = _resolve_log_dir()
    handlers = [logging.StreamHandler()]
    try:
        log_dir.mkdir(parents=True, exist_ok=True)
        handlers.insert(0, logging.FileHandler(log_dir / "daemon.log"))
    except Exception:
        # Fall back to /tmp if the default log dir is not writable.
        try:
            tmp_dir = Path(os.environ.get("TMPDIR", "/tmp")) / "sari"
            tmp_dir.mkdir(parents=True, exist_ok=True)
            handlers.insert(0, logging.FileHandler(tmp_dir / "daemon.log"))
        except Exception:
            pass

    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
        handlers=handlers,
    )

_init_logging()
logger = logging.getLogger("mcp-daemon")

DEFAULT_HOST = DEFAULT_DAEMON_HOST
DEFAULT_PORT = DEFAULT_DAEMON_PORT
PID_FILE = WorkspaceManager.get_global_data_dir() / "daemon.pid"

def _pid_file() -> Path:
    try:
        if isinstance(PID_FILE, Path):
            return PID_FILE
    except Exception:
        pass
    return WorkspaceManager.get_global_data_dir() / "daemon.pid"

class SariDaemon:
    def __init__(self, host: str = None, port: int = None):
        self.host = host or settings.DAEMON_HOST
        self.port = int(port or settings.DAEMON_PORT)
        self.server = None
        self._loop = None
        self._pinned_workspace_root = None
        self._registry = ServerRegistry()
        
        # Use existing boot ID if provided, otherwise generate new
        env_boot_id = (os.environ.get("SARI_BOOT_ID") or "").strip()
        self.boot_id = env_boot_id or uuid.uuid4().hex
        if not env_boot_id:
            os.environ["SARI_BOOT_ID"] = self.boot_id
            
        self._stop_event = threading.Event()
        self._heartbeat_thread = None
        self._idle_since = None
        self._drain_since = None
        self._autostop_no_client_since = None
        self.httpd = None
        self.http_host = None
        self.http_port = None
        self._active_connections = 0
        self._conn_lock = threading.Lock()
        trace("daemon_init", host=self.host, port=self.port)

    def _cleanup_legacy_pid_file(self):
        """Best-effort cleanup for legacy pid file; registry is the SSOT."""
        try:
            pid_file = _pid_file()
            if pid_file.exists():
                pid_file.unlink()
                logger.info("Removed legacy daemon.pid file")
        except Exception as e:
            logger.debug(f"Failed to remove legacy daemon.pid file: {e}")

    def _register_daemon(self):
        try:
            sari_version = settings.VERSION
            host = (self.host or "127.0.0.1").strip()
            port = int(self.port)
            pid = os.getpid()
            self._registry.register_daemon(self.boot_id, host, port, pid, version=sari_version)
            logger.info(f"Registered daemon {self.boot_id} on {host}:{port}")
            trace("daemon_registered", boot_id=self.boot_id, host=host, port=port, pid=pid, version=sari_version)
        except Exception as e:
            logger.error(f"Failed to register daemon: {e}")
            trace("daemon_register_error", error=str(e))

    def _unregister_daemon(self):
        try:
            self._registry.unregister_daemon(self.boot_id)
            logger.info(f"Unregistered daemon {self.boot_id}")
            trace("daemon_unregistered", boot_id=self.boot_id)
        except Exception as e:
            logger.error(f"Failed to unregister daemon: {e}")
            trace("daemon_unregister_error", error=str(e))

    def _autostart_workspace(self) -> None:
        if not settings.DAEMON_AUTOSTART:
            return

        workspace_root = settings.WORKSPACE_ROOT or WorkspaceManager.resolve_workspace_root()

        try:
            from sari.mcp.workspace_registry import Registry
            # Do not pin a permanent reference here; allow autostop when CLI sessions close.
            shared = Registry.get_instance().get_or_create(
                workspace_root,
                persistent=False,
                track_ref=False,
            )
            self._pinned_workspace_root = workspace_root
            self._start_http_gateway(shared)
            logger.info(f"Auto-started workspace session for {workspace_root}")
        except Exception as e:
            logger.error(f"Failed to auto-start workspace session: {e}")
        try:
            # Record workspace mapping even if warm-up failed.
            self._registry.set_workspace(
                workspace_root,
                self.boot_id,
                http_port=self.http_port,
                http_host=self.http_host,
            )
        except Exception as e:
            logger.error(f"Failed to record workspace mapping: {e}")

    def _start_http_gateway(self, shared_state) -> None:
        if self.httpd is not None:
            return
        try:
            from sari.core.http_server import serve_forever

            host = "127.0.0.1"
            port = int(settings.HTTP_API_PORT)
            httpd, actual_port = serve_forever(
                host,
                port,
                shared_state.db,
                shared_state.indexer,
                version=settings.VERSION,
                workspace_root=str(shared_state.workspace_root),
                mcp_server=shared_state.server,
                shared_http_gateway=True,
            )
            self.httpd = httpd
            self.http_host = host
            self.http_port = actual_port
            # Shared gateway runs in-process with daemon; avoid a misleading separate http_pid.
            self._registry.set_daemon_http(self.boot_id, actual_port, http_host=host)
            logger.info(f"Started daemon HTTP gateway on {host}:{actual_port}")
        except Exception as e:
            logger.error(f"Failed to start daemon HTTP gateway: {e}")

    def _inc_active_connections(self) -> None:
        with self._conn_lock:
            self._active_connections += 1

    def _dec_active_connections(self) -> None:
        with self._conn_lock:
            self._active_connections = max(0, self._active_connections - 1)

    def _get_active_connections(self) -> int:
        with self._conn_lock:
            return self._active_connections

    def _start_heartbeat(self) -> None:
        if self._heartbeat_thread and self._heartbeat_thread.is_alive():
            return
        self._heartbeat_thread = threading.Thread(target=self._heartbeat_loop, daemon=True)
        self._heartbeat_thread.start()

    def _heartbeat_loop(self) -> None:
        interval = settings.DAEMON_HEARTBEAT_SEC
        idle_sec = settings.DAEMON_IDLE_SEC
        idle_with_active = settings.DAEMON_IDLE_WITH_ACTIVE
        drain_grace = settings.DAEMON_DRAIN_GRACE_SEC
        autostop_grace = max(1, int(settings.get_int("DAEMON_AUTOSTOP_GRACE_SEC", 60)))
        autostop_enabled = bool(settings.get_bool("DAEMON_AUTOSTOP", True))

        from sari.mcp.workspace_registry import Registry
        workspace_registry = Registry.get_instance()

        while not self._stop_event.is_set():
            try:
                self._registry.touch_daemon(self.boot_id)
                daemon_info = self._registry.get_daemon(self.boot_id) or {}
                draining = bool(daemon_info.get("draining"))
                active_count = workspace_registry.active_count()
                socket_active = self._get_active_connections()
                indexing_active = bool(
                    getattr(
                        workspace_registry,
                        "has_indexing_activity",
                        lambda: False)())
                now = time.time()

                # If no live socket clients remain, reap stale workspace refs to avoid
                # daemon pinning after abnormal client termination.
                if socket_active == 0 and active_count > 0:
                    reap_fn = getattr(workspace_registry, "reap_stale_refs", None)
                    if callable(reap_fn):
                        try:
                            reap_fn(max(1, float(autostop_grace)))
                            active_count = workspace_registry.active_count()
                        except Exception:
                            pass

                # Prefer real socket liveness as a client signal.
                active_count = max(active_count, socket_active)

                if draining:
                    if self._drain_since is None:
                        self._drain_since = now
                    if active_count == 0:
                        self.shutdown()
                        break
                    if drain_grace > 0 and now - self._drain_since >= drain_grace:
                        self.shutdown()
                        break
                else:
                    self._drain_since = None
                    if autostop_enabled:
                        if active_count == 0 and not indexing_active:
                            if self._autostop_no_client_since is None:
                                self._autostop_no_client_since = now
                            elif now - self._autostop_no_client_since >= autostop_grace:
                                self.shutdown()
                                break
                        else:
                            self._autostop_no_client_since = None
                    if idle_sec > 0:
                        if (active_count == 0 or idle_with_active) and not indexing_active:
                            last_activity = workspace_registry.get_last_activity_ts()
                            if last_activity <= 0:
                                last_activity = now
                            if self._idle_since is None:
                                self._idle_since = last_activity
                            # Keep a small grace window to avoid flapping right after startup.
                            if now - last_activity >= (idle_sec + 1.0):
                                self.shutdown()
                                break
                        else:
                            self._idle_since = None
            except Exception as e:
                logger.error(f"Heartbeat failed: {e}")

            self._stop_event.wait(interval)

    async def start_async(self):
        host = (self.host or "127.0.0.1").strip()
        try:
            is_loopback = host.lower() == "localhost" or ipaddress.ip_address(host).is_loopback
        except ValueError:
            is_loopback = host.lower() == "localhost"

        if not is_loopback:
            raise SystemExit(
                f"sari daemon refused to start: host must be loopback only (127.0.0.1/localhost/::1). got={host}. "
                "Remote access is NOT supported for security."
            )

        # Prevent duplicate starts on the same endpoint
        existing = self._registry.resolve_daemon_by_endpoint(host, self.port)
        if existing:
            raise SystemExit(f"sari daemon already running on {host}:{self.port} (PID: {existing['pid']})")

        self._cleanup_legacy_pid_file()
        trace("daemon_start_async", host=host, port=self.port)
        self.server = await asyncio.start_server(
            self.handle_client, self.host, self.port
        )
        os.environ["SARI_DAEMON_HOST"] = host
        os.environ["SARI_DAEMON_PORT"] = str(self.port)
        self._loop = asyncio.get_running_loop()
        self._register_daemon()
        self._autostart_workspace()
        self._start_heartbeat()

        addr = self.server.sockets[0].getsockname()
        logger.info(f"Sari Daemon serving on {addr}")
        trace("daemon_listening", addr=str(addr))

        async with self.server:
            runner = self.server.serve_forever()
            if inspect.isawaitable(runner):
                await runner
            else:
                # Test doubles may provide non-awaitable mocks.
                while not self._stop_event.is_set():
                    await asyncio.sleep(0.1)

    def start(self):
        try:
            asyncio.get_running_loop()
            return self.start_async()
        except RuntimeError:
            try:
                asyncio.run(self.start_async())
            except asyncio.CancelledError:
                pass
            return None

    async def handle_client(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter):
        addr = writer.get_extra_info('peername')
        logger.info(f"Accepted connection from {addr}")
        trace("daemon_client_accepted", addr=str(addr))
        self._inc_active_connections()
        try:
            session = Session(reader, writer)
            await session.handle_connection()
        finally:
            self._dec_active_connections()
            logger.info(f"Closed connection from {addr}")
            trace("daemon_client_closed", addr=str(addr))

    def stop(self):
        """Public stop method for external control (tests, etc.)."""
        self.shutdown()

    def shutdown(self):
        if self._stop_event.is_set():
            return
        self._stop_event.set()

        logger.info("Initiating graceful shutdown...")
        trace("daemon_shutdown_start")

        # 1. Stop Server Loop
        if self.server:
            try:
                self.server.close()
                if self._loop and self._loop.is_running():
                    # Ensure close callback is processed
                    pass
            except Exception as e:
                logger.debug(f"Error during server close: {e}")

        # 2. Shutdown all workspaces (Stops Indexers and ProcessPools)
        from sari.mcp.workspace_registry import Registry
        try:
            logger.info("Shutting down workspace registry and indexers...")
            Registry.get_instance().shutdown_all()
        except Exception as e:
            logger.error(f"Error shutting down registry: {e}")

        if self.httpd is not None:
            try:
                self.httpd.shutdown()
                self.httpd.server_close()
            except Exception as e:
                logger.debug(f"Error shutting down HTTP gateway: {e}")

        # 3. Kill any remaining children (Safety Net)
        try:
            import multiprocessing
            for child in multiprocessing.active_children():
                logger.warning(f"Terminating lingering child process: {child.pid}")
                child.terminate()
        except Exception:
            pass

        self._unregister_daemon()
        self._cleanup_legacy_pid_file()
        logger.info("Shutdown sequence complete.")
        trace("daemon_shutdown_complete")

async def main():
    daemon = SariDaemon()

    # Handle signals
    loop = asyncio.get_running_loop()
    stop_event = asyncio.Event()

    def _handle_signal():
        stop_event.set()

    try:
        loop.add_signal_handler(signal.SIGTERM, _handle_signal)
        loop.add_signal_handler(signal.SIGINT, _handle_signal)
    except Exception:
        # Some environments (tests/platforms) may not support signal handlers.
        pass

    daemon_task = asyncio.create_task(daemon.start_async())
    stop_task = asyncio.create_task(stop_event.wait())

    logger.info("Daemon started. Press Ctrl+C to stop.")

    try:
        done, _ = await asyncio.wait(
            {daemon_task, stop_task},
            return_when=asyncio.FIRST_COMPLETED,
        )
        if daemon_task in done:
            # Propagate startup/runtime errors immediately.
            exc = daemon_task.exception()
            if exc is not None:
                raise exc
    finally:
        stop_task.cancel()
        logger.info("Stopping daemon...")
        daemon.shutdown()
        if not daemon_task.done():
            daemon_task.cancel()
        try:
            await daemon_task
        except asyncio.CancelledError:
            pass
        logger.info("Daemon stopped.")

if __name__ == "__main__":
    asyncio.run(main())
