"""Common utilities for MCP server.

This module provides shared utility functions used across the MCP server,
including file polling and other async helpers.
"""

from __future__ import annotations

import asyncio
import enum
import logging
import time
from dataclasses import dataclass
from pathlib import Path
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from .client import FastAPIClient

__all__ = [
    "CellBuildResult",
    "CellBuildStatus",
    "poll_cell_build_status",
    "wait_for_build_completion",
    "wait_for_build_completion_via_api",
    "wait_for_gds_file",
]

logger = logging.getLogger(__name__)


async def wait_for_gds_file(
    gds_path: Path,
    timeout: float = 30.0,
    poll_interval: float = 0.5,
    min_age: float = 0.1,
) -> bool:
    """Poll for GDS file to be created with timeout.

    This function waits for a GDS file to exist and be stable (not actively
    being written to) before returning. Useful for waiting on build operations.

    Args:
        gds_path: Path to the expected GDS file
        timeout: Maximum time to wait in seconds (default: 30.0)
        poll_interval: Time between polls in seconds (default: 0.5)
        min_age: Minimum file age in seconds to ensure write is complete (default: 0.1)

    Returns:
        True if file exists and is ready, False if timeout reached
    """
    start_time = time.monotonic()
    logger.debug("Waiting for GDS file: %s (timeout: %.1fs)", gds_path, timeout)

    while time.monotonic() - start_time < timeout:
        if gds_path.exists():
            file_age = time.time() - gds_path.stat().st_mtime
            if file_age >= min_age:
                logger.debug("GDS file ready: %s (age: %.2fs)", gds_path, file_age)
                return True
            logger.debug(
                "GDS file exists but too new (age: %.2fs), waiting...", file_age
            )

        await asyncio.sleep(poll_interval)

    logger.warning("Timeout waiting for GDS file: %s", gds_path)
    return False


async def wait_for_build_completion(
    cell_names: list[str],
    project: str | None = None,
    timeout: float = 30.0,
    poll_interval: float = 0.5,
) -> dict[str, bool]:
    """Wait for all cells to finish building by polling for GDS files.

    Resolves the project path from the server registry and waits for all
    expected GDS files to appear on disk.

    Args:
        cell_names: List of cell names that were built
        project: Optional project name for registry lookup
        timeout: Maximum time to wait per cell in seconds (default: 30.0)
        poll_interval: Time between polls in seconds (default: 0.5)

    Returns:
        Dict mapping cell name to bool (True if ready, False if timed out).
        Returns empty dict if no servers found in registry.
    """
    from .registry import ServerRegistry

    if not cell_names:
        return {}

    registry = ServerRegistry()
    server_info = None

    if project:
        server_info = registry.get_server_by_project(project)
    if not server_info:
        servers = registry.list_servers()
        if servers:
            server_info = servers[0]

    if not server_info:
        logger.debug("No server found in registry, skipping build completion wait")
        return {}

    project_path = Path(server_info.project_path)
    gds_dir = project_path / "build" / "gds"

    async def _wait_for_cell(cell_name: str) -> tuple[str, bool]:
        gds_path = gds_dir / f"{cell_name}.gds"
        ready = await wait_for_gds_file(
            gds_path, timeout=timeout, poll_interval=poll_interval
        )
        return cell_name, ready

    results = await asyncio.gather(*[_wait_for_cell(name) for name in cell_names])
    return dict(results)


class CellBuildStatus(enum.Enum):
    """Status of a cell build."""

    READY = "ready"
    FAILED = "failed"
    TIMED_OUT = "timed_out"


@dataclass
class CellBuildResult:
    """Result of polling a cell's build status."""

    status: CellBuildStatus
    error_message: str | None = None


async def poll_cell_build_status(
    client: FastAPIClient,
    cell_name: str,
    project: str | None = None,
    timeout: float = 30.0,
    poll_interval: float = 0.5,
) -> CellBuildResult:
    """Poll the build status API for a single cell.

    Args:
        client: FastAPI client for HTTP requests
        cell_name: Name of the cell to poll
        project: Optional project name for routing
        timeout: Maximum time to wait in seconds
        poll_interval: Time between polls in seconds

    Returns:
        CellBuildResult with status and optional error message
    """
    import httpx

    start_time = time.monotonic()
    logger.debug(
        "Polling build status for cell: %s (timeout: %.1fs)", cell_name, timeout
    )

    while time.monotonic() - start_time < timeout:
        try:
            response = await client.request(
                method="GET",
                path=f"/api/build-status/{cell_name}",
                project=project,
            )

            status = response.get("status") if isinstance(response, dict) else None

            if status == "success":
                logger.debug("Cell %s build succeeded", cell_name)
                return CellBuildResult(status=CellBuildStatus.READY)

            if status == "failed":
                error = (
                    response.get("error", "Unknown error")
                    if isinstance(response, dict)
                    else "Unknown error"
                )
                logger.debug("Cell %s build failed: %s", cell_name, error)
                return CellBuildResult(
                    status=CellBuildStatus.FAILED, error_message=error
                )

            # "building" or "unknown" — keep polling
            logger.debug("Cell %s status: %s, continuing to poll...", cell_name, status)

        except httpx.HTTPStatusError as e:
            if e.response.status_code == 404:
                # Race condition: build may not have started yet
                logger.debug(
                    "Cell %s not found yet (404), continuing to poll...", cell_name
                )
            else:
                logger.warning("HTTP error polling cell %s: %s", cell_name, e)
        except Exception:
            logger.warning(
                "Error polling build status for cell %s", cell_name, exc_info=True
            )

        await asyncio.sleep(poll_interval)

    logger.warning("Timeout polling build status for cell: %s", cell_name)
    return CellBuildResult(status=CellBuildStatus.TIMED_OUT)


async def wait_for_build_completion_via_api(
    client: FastAPIClient,
    cell_names: list[str],
    project: str | None = None,
    timeout: float = 30.0,
    poll_interval: float = 0.5,
) -> dict[str, CellBuildResult]:
    """Wait for all cells to finish building by polling the build status API.

    Uses asyncio.gather() to poll all cells in parallel.

    Args:
        client: FastAPI client for HTTP requests
        cell_names: List of cell names to poll
        project: Optional project name for routing
        timeout: Maximum time to wait per cell in seconds
        poll_interval: Time between polls in seconds

    Returns:
        Dict mapping cell name to CellBuildResult
    """
    if not cell_names:
        return {}

    async def _poll_cell(name: str) -> tuple[str, CellBuildResult]:
        result = await poll_cell_build_status(
            client, name, project=project, timeout=timeout, poll_interval=poll_interval
        )
        return name, result

    results = await asyncio.gather(*[_poll_cell(name) for name in cell_names])
    return dict(results)
