import { z } from "zod";

/** Per-request timeout (paginated search may run several requests). */
const SEARCH_PAGE_TIMEOUT_MS = 60_000;

/** Datadog allows up to 1000 logs per page for POST /api/v2/logs/events/search. */
const MAX_PAGE_SIZE = 1000;

/** Hard cap on total logs returned in one API call chain (safety). */
const MAX_TOTAL_LOGS_CAP = 50_000;

/** Avoid unbounded pagination if cursor never clears. */
const MAX_PAGES = 100;

/** Max retries on 429 Rate Limit responses. */
const MAX_RATE_LIMIT_RETRIES = 3;

/** Fallback wait time (ms) when Datadog doesn't send a Retry-After / x-ratelimit-reset header. */
const RATE_LIMIT_FALLBACK_WAIT_MS = 10_000;

/** Response shape for POST /api/v2/logs/events/search (subset). */
const logEventAttributesSchema = z.object({
  message: z.string().optional(),
  timestamp: z.string().optional(),
  status: z.string().optional(),
  service: z.string().optional(),
  host: z.string().optional(),
  attributes: z.record(z.string(), z.unknown()).optional(),
});

const logEventSchema = z.object({
  id: z.string().optional(),
  type: z.string().optional(),
  attributes: logEventAttributesSchema.optional(),
});

const searchResponseSchema = z.object({
  data: z.array(logEventSchema).optional(),
  meta: z
    .object({
      page: z
        .object({
          after: z.string().optional(),
        })
        .optional(),
    })
    .optional(),
});

export interface DatadogLogSummary {
  id: string | null;
  timestamp: string | null;
  service: string | null;
  status: string | null;
  host: string | null;
  source: string | null;
  message: string | null;
}

export interface DatadogLogsSearchResult {
  logs: DatadogLogSummary[];
  total: number;
  rawMeta: unknown;
  /** True when more logs matched the query but `limit` stopped collection. */
  truncated?: boolean;
}

export interface DatadogLogsSearchParams {
  query: string;
  from: string;
  to: string;
  /** Max logs to return across all pages (capped internally). */
  limit: number;
}

export class DatadogApiError extends Error {
  constructor(
    message: string,
    public readonly status: number,
    public readonly responseBody: string,
  ) {
    super(message);
    this.name = "DatadogApiError";
  }
}

function extractSource(attrs: Record<string, unknown> | undefined): string | null {
  if (!attrs) return null;
  const src = attrs["ddsource"];
  if (typeof src === "string") return src;
  return null;
}

function mapEventsToSummaries(data: z.infer<typeof logEventSchema>[]): DatadogLogSummary[] {
  return data.map((item) => {
    const a = item.attributes;
    const nested = a?.attributes as Record<string, unknown> | undefined;
    const hostFromNested = nested && typeof nested["host"] === "string" ? nested["host"] : null;
    return {
      id: item.id ?? null,
      timestamp: a?.timestamp ?? null,
      service: a?.service ?? null,
      status: a?.status ?? null,
      host: typeof a?.host === "string" ? a.host : hostFromNested,
      source: extractSource(nested),
      message: a?.message ?? null,
    };
  });
}

export class DatadogLogsClient {
  private readonly baseUrl: string;

  constructor(
    private readonly apiKey: string,
    private readonly appKey: string,
    site: string,
  ) {
    this.baseUrl = `https://api.${site}`;
  }

  /**
   * Searches logs in [from, to], following Datadog cursor pagination until
   * `limit` logs are collected, no more pages exist, or safety caps hit.
   */
  async searchLogs(params: DatadogLogsSearchParams): Promise<DatadogLogsSearchResult> {
    const url = `${this.baseUrl}/api/v2/logs/events/search`;
    const maxTotal = Math.min(Math.max(1, params.limit), MAX_TOTAL_LOGS_CAP);
    const allLogs: DatadogLogSummary[] = [];
    let cursor: string | undefined;
    let lastMeta: unknown;
    let truncated = false;
    let lastNextCursor: string | undefined;

    for (let page = 0; page < MAX_PAGES && allLogs.length < maxTotal; page++) {
      const remaining = maxTotal - allLogs.length;
      const pageLimit = Math.min(MAX_PAGE_SIZE, remaining);

      const pageBody: Record<string, unknown> = {
        filter: {
          query: params.query,
          from: params.from,
          to: params.to,
        },
        page: {
          limit: pageLimit,
          ...(cursor ? { cursor } : {}),
        },
        sort: "-timestamp",
      };

      let response: Response;
      let text: string;

      // Retry on 429 up to MAX_RATE_LIMIT_RETRIES times, honouring Retry-After.
      for (let attempt = 0; ; attempt++) {
        const controller = new AbortController();
        const timeoutId = setTimeout(() => controller.abort(), SEARCH_PAGE_TIMEOUT_MS);

        try {
          response = await fetch(url, {
            method: "POST",
            headers: {
              "Content-Type": "application/json",
              Accept: "application/json",
              "DD-API-KEY": this.apiKey,
              "DD-APPLICATION-KEY": this.appKey,
            },
            body: JSON.stringify(pageBody),
            signal: controller.signal,
          });
        } catch (err) {
          clearTimeout(timeoutId);
          const message = err instanceof Error ? err.message : String(err);
          if (err instanceof Error && err.name === "AbortError") {
            throw new DatadogApiError("Datadog logs search timed out", 504, message);
          }
          throw new DatadogApiError(`Datadog request failed: ${message}`, 502, message);
        }
        clearTimeout(timeoutId);

        if (response.status === 429 && attempt < MAX_RATE_LIMIT_RETRIES) {
          const retryAfterHeader =
            response.headers.get("x-ratelimit-reset") ?? response.headers.get("retry-after");
          const waitSeconds = retryAfterHeader ? parseInt(retryAfterHeader, 10) : NaN;
          const waitMs = Number.isNaN(waitSeconds)
            ? RATE_LIMIT_FALLBACK_WAIT_MS
            : Math.min(waitSeconds * 1000, 60_000);
          await new Promise((r) => setTimeout(r, waitMs));
          continue;
        }

        text = await response.text();
        if (!response.ok) {
          throw new DatadogApiError(
            `Datadog API error (${response.status})`,
            response.status,
            text.slice(0, 2000),
          );
        }
        break;
      }

      let json: unknown;
      try {
        json = JSON.parse(text) as unknown;
      } catch {
        throw new DatadogApiError("Invalid JSON from Datadog", 502, text.slice(0, 500));
      }

      const parsed = searchResponseSchema.safeParse(json);
      if (!parsed.success) {
        throw new DatadogApiError("Unexpected Datadog response shape", 502, text.slice(0, 500));
      }

      const rows = parsed.data.data ?? [];
      lastMeta = parsed.data.meta ?? null;
      const nextCursor = parsed.data.meta?.page?.after;
      lastNextCursor = nextCursor;

      allLogs.push(...mapEventsToSummaries(rows));

      if (allLogs.length >= maxTotal) {
        truncated = Boolean(nextCursor);
        break;
      }

      if (!nextCursor || rows.length === 0) {
        break;
      }

      cursor = nextCursor;
    }

    // More pages existed but we stopped at MAX_PAGES (loop never consumed `cursor` for next request).
    if (!truncated && lastNextCursor) {
      truncated = true;
    }

    return {
      logs: allLogs,
      total: allLogs.length,
      rawMeta: lastMeta,
      ...(truncated ? { truncated: true } : {}),
    };
  }
}
