"""MCP tool registrations for log querying via Loki."""

from __future__ import annotations

import json
import os
import urllib.error
import urllib.parse
import urllib.request
from datetime import datetime, timedelta, timezone

from mcp.server.fastmcp import Context


def _loki_url() -> str:
    return os.environ.get("PEON_LOKI_URL", "http://loki:3100").rstrip("/")


def _iso_to_ns(ts: str) -> int:
    """Convert an ISO 8601 timestamp string to nanoseconds since the Unix epoch."""
    # Handle common ISO formats: with/without Z or +00:00 suffix
    ts = ts.strip()
    for fmt in (
        "%Y-%m-%dT%H:%M:%SZ",
        "%Y-%m-%dT%H:%M:%S.%fZ",
        "%Y-%m-%dT%H:%M:%S+00:00",
        "%Y-%m-%dT%H:%M:%S",
        "%Y-%m-%d %H:%M:%S",
    ):
        try:
            dt = datetime.strptime(ts, fmt).replace(tzinfo=timezone.utc)
            return int(dt.timestamp() * 1_000_000_000)
        except ValueError:
            continue
    raise ValueError(f"Cannot parse timestamp: {ts!r}")


def register_tools(mcp, api_fn):
    @mcp.tool()
    async def query_logs(
        query: str,
        start: str = "",
        end: str = "",
        limit: int = 100,
        ctx: Context = None,
    ) -> list[dict]:
        """Query production logs via the Loki HTTP API.

        Sends a LogQL query to the Loki instance configured by the
        ``PEON_LOKI_URL`` environment variable (default: ``http://loki:3100``).

        Args:
            query: LogQL query string, e.g. ``'{app="peon-api"} |= "ERROR"'``.
            start: ISO 8601 start timestamp (default: 15 minutes ago).
            end: ISO 8601 end timestamp (default: now).
            limit: Maximum number of log lines to return (default 100).

        Returns:
            A list of dicts with keys ``timestamp`` (ISO 8601), ``labels``
            (dict), and ``line`` (the raw log line).
        """
        now = datetime.now(timezone.utc)
        start_ns = _iso_to_ns(start) if start else int((now - timedelta(minutes=15)).timestamp() * 1_000_000_000)
        end_ns = _iso_to_ns(end) if end else int(now.timestamp() * 1_000_000_000)

        params = urllib.parse.urlencode(
            {
                "query": query,
                "start": str(start_ns),
                "end": str(end_ns),
                "limit": str(limit),
                "direction": "backward",
            }
        )
        url = f"{_loki_url()}/loki/api/v1/query_range?{params}"

        req = urllib.request.Request(url, headers={"Accept": "application/json"})
        try:
            with urllib.request.urlopen(req, timeout=15) as resp:
                data = json.loads(resp.read())
        except urllib.error.HTTPError as exc:
            body = exc.read().decode(errors="replace")
            raise RuntimeError(f"Loki HTTP {exc.code}: {body}") from exc
        except (urllib.error.URLError, OSError) as exc:
            raise RuntimeError(f"Loki unreachable: {exc}") from exc

        # Loki response: {"data": {"result": [{"stream": {...}, "values": [[ns, line], ...]}, ...]}}
        entries: list[dict] = []
        for stream in data.get("data", {}).get("result", []):
            labels = stream.get("stream", {})
            for ns_str, line in stream.get("values", []):
                ts = datetime.fromtimestamp(int(ns_str) / 1_000_000_000, tz=timezone.utc)
                entries.append(
                    {
                        "timestamp": ts.strftime("%Y-%m-%dT%H:%M:%S.%f")[:-3] + "Z",
                        "labels": labels,
                        "line": line,
                    }
                )

        return entries
