# ========= Copyright 2026 @ Strukto.AI All Rights Reserved. =========
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ========= Copyright 2026 @ Strukto.AI All Rights Reserved. =========

import asyncio
import logging
import os
import signal
import time
from contextlib import asynccontextmanager
from pathlib import Path

from fastapi import FastAPI

from mirage.server.jobs import JobTable
from mirage.server.persist import restore_all, snapshot_all
from mirage.server.registry import WorkspaceRegistry
from mirage.server.routers import execute, health, jobs, sessions, workspaces

logger = logging.getLogger(__name__)


def _pid_file_path() -> Path:
    return Path.home() / ".mirage" / "daemon.pid"


def _write_pid_file() -> None:
    p = _pid_file_path()
    p.parent.mkdir(parents=True, exist_ok=True)
    p.write_text(str(os.getpid()))


def _remove_pid_file() -> None:
    try:
        _pid_file_path().unlink(missing_ok=True)
    except OSError:
        logger.debug("could not remove pid file %s", _pid_file_path())


async def _watch_exit(exit_event: asyncio.Event) -> None:
    """Send SIGTERM to self when ``exit_event`` is set.

    Lets uvicorn handle its own graceful shutdown sequence.
    """
    try:
        await exit_event.wait()
    except asyncio.CancelledError:
        return
    logger.info("exit event tripped; sending SIGTERM to self")
    os.kill(os.getpid(), signal.SIGTERM)


@asynccontextmanager
async def _lifespan(app: FastAPI):
    _write_pid_file()
    exit_task = asyncio.create_task(_watch_exit(app.state.exit_event))
    persist_dir = getattr(app.state, "persist_dir", None)
    if persist_dir is not None:
        try:
            restored, skipped = restore_all(app.state.registry,
                                            Path(persist_dir))
            logger.info("restored %d workspaces (%d skipped) from %s",
                        restored, skipped, persist_dir)
        except Exception:
            logger.exception("restore_all failed; starting empty")
    try:
        yield
    finally:
        exit_task.cancel()
        if persist_dir is not None:
            try:
                saved = await snapshot_all(app.state.registry,
                                           Path(persist_dir))
                logger.info("snapshotted %d workspaces to %s", saved,
                            persist_dir)
            except Exception:
                logger.exception("snapshot_all failed")
        await app.state.registry.close_all()
        _remove_pid_file()


def build_app(idle_grace_seconds: float = 30.0,
              exit_event: asyncio.Event | None = None,
              persist_dir: str | Path | None = None) -> FastAPI:
    """Construct a daemon FastAPI app.

    The workspace registry is created eagerly so the app is usable
    even without ASGI lifespan events firing (e.g. inside an
    ``httpx.ASGITransport`` test client). The lifespan handles
    persist-on-shutdown / restore-on-start when ``persist_dir`` is
    supplied.

    Args:
        idle_grace_seconds (float): seconds to wait after the last
            workspace is removed before signalling shutdown.
        exit_event (asyncio.Event | None): event the registry trips
            when the idle timer fires. The runner of this app should
            await it and shut uvicorn down. Defaults to a fresh event.
        persist_dir (str | Path | None): directory for snapshot /
            restore. When set, the daemon snapshots every workspace
            on graceful shutdown and rehydrates them on next start.

    Returns:
        FastAPI: configured app with all routers mounted.
    """
    app = FastAPI(title="Mirage daemon", version="0.1", lifespan=_lifespan)
    app.state.started_at = time.time()
    app.state.exit_event = exit_event or asyncio.Event()
    app.state.registry = WorkspaceRegistry(
        idle_grace_seconds=idle_grace_seconds,
        exit_event=app.state.exit_event,
    )
    app.state.jobs = JobTable()
    app.state.persist_dir = persist_dir
    app.include_router(workspaces.router)
    app.include_router(sessions.router)
    app.include_router(execute.router)
    app.include_router(jobs.router)
    app.include_router(health.router)
    return app
