# Copyright 2026 Marimo. All rights reserved.
from __future__ import annotations

import asyncio
from functools import lru_cache
from typing import TYPE_CHECKING, Any, Literal, TypedDict

from starlette.authentication import requires
from starlette.responses import JSONResponse, PlainTextResponse

from marimo import _loggers
from marimo._dependencies.dependencies import DependencyManager
from marimo._server.api.deps import AppState
from marimo._server.router import APIRouter
from marimo._utils.diagnostics import get_system_info
from marimo._utils.health import (
    get_cgroup_cpu_percent,
    get_cgroup_mem_stats,
    get_node_version,
    get_python_version,
    get_required_modules_list,
)
from marimo._version import __version__

if TYPE_CHECKING:
    from starlette.requests import Request

LOGGER = _loggers.marimo_logger()

# Router for health/status endpoints
router = APIRouter()


def health_check(request: Request) -> JSONResponse:
    del request  # Unused
    return JSONResponse({"status": "healthy"})


# Multiple health endpoints to make it easier on the consumer
router.add_route("/health", health_check, methods=["GET"])
router.add_route("/healthz", health_check, methods=["GET"])


@router.get("/api/status")
@requires("edit")
async def status(request: Request) -> JSONResponse:
    """
    responses:
        200:
            description: Get the status of the application
            content:
                application/json:
                    schema:
                        type: object
                        properties:
                            status:
                                type: string
                            filenames:
                                type: array
                                items:
                                    type: string
                            mode:
                                type: string
                            sessions:
                                type: integer
                            version:
                                type: string
                            requirements:
                                type: array
                                items:
                                    type: string
                            node_version:
                                type: string
                            lsp_running:
                                type: boolean
    """
    app_state = AppState(request)
    files = [
        session.app_file_manager.filename or "__new__"
        for session in app_state.session_manager.sessions.values()
    ]
    return JSONResponse(
        {
            "status": "healthy",
            "filenames": files,
            "mode": app_state.mode,
            "sessions": len(app_state.session_manager.sessions),
            "version": __version__,
            "python_version": get_python_version(),
            "requirements": get_required_modules_list(),
            "node_version": get_node_version(),
            "lsp_running": app_state.session_manager.lsp_server.is_running(),
        }
    )


class SessionInfo(TypedDict):
    filename: str | None
    path: str | None


@router.get("/api/sessions", include_in_schema=False)
@requires("edit")
async def list_sessions(request: Request) -> JSONResponse:
    """List active session IDs and their notebook paths."""

    state = AppState(request)

    sessions = {
        session_id: SessionInfo(
            filename=session.app_file_manager.filename,
            path=session.app_file_manager.path,
        )
        for session_id, session in state.session_manager.sessions.items()
    }
    return JSONResponse(sessions)


@router.get("/api/version")
@requires("read")
async def version(request: Request) -> PlainTextResponse:
    """
    responses:
        200:
            description: Get the version of the application
            content:
                text/plain:
                    schema:
                        type: string
    """
    del request  # Unused
    return PlainTextResponse(__version__)


@router.get("/api/environment")
@requires("edit")
async def environment(request: Request) -> JSONResponse:
    """
    responses:
        200:
            description: Environment information for issue reporting
            content:
                application/json:
                    schema:
                        type: object
                        properties:
                            marimo:
                                type: string
                            editable:
                                type: boolean
                            location:
                                type: string
                            OS:
                                type: string
                            OS Version:
                                type: string
                            Processor:
                                type: string
                            Python Version:
                                type: string
                            Locale:
                                type: string
                            Binaries:
                                type: object
                                additionalProperties:
                                    type: string
                            Dependencies:
                                type: object
                                additionalProperties:
                                    type: string
                            Optional Dependencies:
                                type: object
                                additionalProperties:
                                    type: string
                            Experimental Flags:
                                type: object
                                additionalProperties: {}
                        required:
                            - marimo
                            - editable
                            - location
                            - OS
                            - OS Version
                            - Processor
                            - Python Version
                            - Locale
                            - Binaries
                            - Dependencies
                            - Optional Dependencies
                            - Experimental Flags
    """
    del request  # Unused
    # Collect off the event loop; the install-location and dependency scans
    # touch the filesystem and can block.
    info = await asyncio.to_thread(get_system_info, redact_home=True)
    return JSONResponse(info)


@router.get("/api/usage")
@requires("edit")
async def usage(request: Request) -> JSONResponse:
    """
    responses:
        200:
            description: Get the current memory and CPU usage of the application
            content:
                application/json:
                    schema:
                        type: object
                        properties:
                            memory:
                                type: object
                                properties:
                                    total:
                                        type: integer
                                    available:
                                        type: integer
                                    percent:
                                        type: number
                                    used:
                                        type: integer
                                    free:
                                        type: integer
                                    has_cgroup_mem_limit:
                                        type: boolean
                                required:
                                    - total
                                    - available
                                    - percent
                                    - used
                                    - free
                                    - has_cgroup_mem_limit
                            server:
                                type: object
                                properties:
                                    memory:
                                        type: integer
                                required:
                                    - memory
                            kernel:
                                type: object
                                properties:
                                    memory:
                                        type: integer
                            cpu:
                                type: object
                                properties:
                                    percent:
                                        type: number
                                required:
                                    - percent
                            gpu:
                                type: array
                                items:
                                    type: object
                                    properties:
                                        index:
                                            type: integer
                                        name:
                                            type: string
                                        memory:
                                            type: object
                                            properties:
                                                total:
                                                    type: integer
                                                used:
                                                    type: integer
                                                free:
                                                    type: integer
                                                percent:
                                                    type: number
                                            required:
                                                - total
                                                - used
                                                - free
                                                - percent
                                    required:
                                        - index
                                        - memory
                                        - name
                        required:
                            - memory
                            - cpu

    """

    try:
        import psutil

        has_psutil = True
    except ImportError:
        has_psutil = False

    if cgroup_mem_stats := get_cgroup_mem_stats():
        memory_stats = {
            "has_cgroup_mem_limit": True,
            **cgroup_mem_stats,
        }
    elif has_psutil:
        memory = psutil.virtual_memory()
        memory_stats = {
            "total": memory.total,
            "available": memory.available,
            "percent": memory.percent,
            "used": memory.used,
            "free": memory.free,
            "has_cgroup_mem_limit": False,
        }
    else:
        memory_stats = {
            "total": None,
            "available": None,
            "percent": None,
            "used": None,
            "free": None,
            "has_cgroup_mem_limit": False,
        }

    cpu = get_cgroup_cpu_percent()
    if cpu is None and has_psutil:
        # interval=None is nonblocking; first call returns meaningless value
        # subsequent calls return delta since last call
        cpu = psutil.cpu_percent(interval=None)

    # Collect kernel PIDs first so we can exclude them from server memory
    kernel_memory: int | None = None
    kernel_pids: set[int] = set()
    server_memory: int | None = None
    session = AppState(request).get_current_session()
    if has_psutil:
        try:
            if session and (pid := session.kernel_pid()) is not None:
                kernel_process = psutil.Process(pid)
                kernel_pids.add(kernel_process.pid)
                kernel_memory = kernel_process.memory_info().rss
                kernel_children = kernel_process.children(recursive=True)
                for child in kernel_children:
                    kernel_pids.add(child.pid)
                    try:
                        kernel_memory += child.memory_info().rss
                    except psutil.NoSuchProcess:
                        pass
        except psutil.ZombieProcess:
            LOGGER.warning("Kernel process is a zombie")

        # Server memory (excluding kernel processes to avoid double-counting)
        main_process = psutil.Process()
        server_memory = main_process.memory_info().rss
        children = main_process.children(recursive=True)
        for child in children:
            if child.pid in kernel_pids:
                continue
            try:
                server_memory += child.memory_info().rss
            except psutil.NoSuchProcess:
                pass

    # GPU stats
    gpu_stats: list[dict[str, Any]] = []
    gpu_available = _is_gpu_available()
    if gpu_available == "nvidia":
        gpu_stats = _parse_nvidia_smi_stats()
    elif gpu_available == "rocm":
        gpu_stats = _parse_rocm_smi_stats()

    return JSONResponse(
        {
            # computer memory or container memory
            "memory": memory_stats,
            # marimo server
            "server": {
                "memory": server_memory,
            },
            # marimo kernel (for the given session)
            "kernel": {
                "memory": kernel_memory,
            },
            "cpu": {
                "percent": cpu,
            },
            "gpu": gpu_stats,
        }
    )


@router.get("/api/status/connections")
@requires("read")
async def connections(request: Request) -> JSONResponse:
    """
    responses:
        200:
            description: Get the number of active websocket connections
            content:
                application/json:
                    schema:
                        type: object
                        properties:
                            active:
                                type: integer
    """
    app_state = AppState(request)
    return JSONResponse(
        {"active": app_state.session_manager.get_active_connection_count()}
    )


_NVIDIA_GPU_STATS_CMD = [
    "nvidia-smi",
    "--query-gpu=index,name,memory.total,memory.used,memory.free",
    "--format=csv,noheader,nounits",
]

_AMD_GPU_STATS_CMD = [
    "rocm-smi",
    "--showproductname",
    "--showuse",
    "--showmeminfo",
    "vram",
    "--json",
]


@lru_cache(maxsize=1)
def _is_gpu_available() -> Literal["nvidia", "rocm", False]:
    import subprocess

    if DependencyManager.which("nvidia-smi"):
        try:
            subprocess.run(
                _NVIDIA_GPU_STATS_CMD,
                capture_output=True,
                text=True,
                check=True,
            )
            return "nvidia"
        except (subprocess.SubprocessError, FileNotFoundError):
            pass

    if DependencyManager.which("rocm-smi"):
        try:
            subprocess.run(
                _AMD_GPU_STATS_CMD,
                capture_output=True,
                text=True,
                check=True,
            )
            return "rocm"
        except (subprocess.SubprocessError, FileNotFoundError):
            pass

    LOGGER.warning(
        "Neither nvidia-smi nor rocm-smi succeeded, cannot get GPU stats"
    )
    return False


def _parse_nvidia_smi_stats() -> list[dict[str, Any]]:
    import subprocess

    try:
        result = subprocess.run(
            _NVIDIA_GPU_STATS_CMD,
            capture_output=True,
            text=True,
            check=True,
        )
    except (subprocess.SubprocessError, FileNotFoundError) as e:
        LOGGER.warning("Failed to extract Nvidia GPU stats: %s", e)
        return []
    stats: list[dict[str, Any]] = []
    for line in result.stdout.strip().split("\n"):
        if not line:
            continue
        index_str, name, total_str, used_str, free_str = line.split(", ")
        # This is what you get on a DGX Spark
        if total_str == "[N/A]":
            total_str = "0"
        if used_str == "[N/A]":
            used_str = "0"
        if free_str == "[N/A]":
            free_str = "0"
        total = int(total_str) * 1024 * 1024  # Convert MB to bytes
        used = int(used_str) * 1024 * 1024
        free = int(free_str) * 1024 * 1024
        stats.append(
            {
                "index": int(index_str),
                "name": name.strip(),
                "memory": {
                    "total": total,
                    "used": used,
                    "free": free,
                    "percent": (used / total) * 100 if total > 0 else 0,
                },
            }
        )
    return stats


def _parse_rocm_smi_stats() -> list[dict[str, Any]]:
    import json
    import subprocess

    try:
        result = subprocess.run(
            _AMD_GPU_STATS_CMD,
            capture_output=True,
            text=True,
            check=True,
        )
    except (subprocess.SubprocessError, FileNotFoundError) as e:
        LOGGER.warning("Failed to extract AMD GPU stats: %s", e)
        return []

    stdout = result.stdout.strip()
    # Find the JSON object (rocm-smi may emit WARNING lines before the JSON)
    json_start = stdout.find("{")
    if json_start == -1:
        LOGGER.warning("No JSON found in rocm-smi output")
        return []
    try:
        data: dict[str, dict[str, str]] = json.loads(stdout[json_start:])
    except json.JSONDecodeError as e:
        LOGGER.warning("Failed to parse rocm-smi JSON: %s", e)
        return []

    stats: list[dict[str, Any]] = []
    for device_key, device_info in data.items():
        if not device_key.startswith("card"):
            continue
        index = int(device_key.replace("card", ""))
        name = device_info.get("Card Series", "AMD GPU")
        total_str = device_info.get("VRAM Total Memory (B)", "0")
        used_str = device_info.get("VRAM Total Used Memory (B)", "0")
        total = int(total_str) if total_str not in ("N/A", "[N/A]", "") else 0
        used = int(used_str) if used_str not in ("N/A", "[N/A]", "") else 0
        total = max(total, used)
        free = total - used
        stats.append(
            {
                "index": index,
                "name": name.strip(),
                "memory": {
                    "total": total,
                    "used": used,
                    "free": free,
                    "percent": (used / total) * 100 if total > 0 else 0,
                },
            }
        )
    return stats
