import type { FormatHandler, ParsedRequest } from "../handler";
import type { CapturedLog, TokenUsage } from "../../schemas";
import { AnthropicRequestSchema, AnthropicResponseSchema } from "./schemas";
import { extractAnthropicStream } from "./stream";

export const AnthropicFormatHandler: FormatHandler = {
  format: "anthropic",

  parseRequest(rawBody: string, headers?: Headers): ParsedRequest | null {
    try {
      const json: unknown = JSON.parse(rawBody);
      const result = AnthropicRequestSchema.safeParse(json);
      if (result.success) {
        return {
          model: result.data.model,
          sessionId: result.data.metadata?.user_id ?? headers?.get("x-session-affinity") ?? null,
        };
      }
      return null;
    } catch {
      return null;
    }
  },

  extractTokens(responseBody: string): TokenUsage {
    try {
      const json: unknown = JSON.parse(responseBody);
      const result = AnthropicResponseSchema.safeParse(json);
      if (result.success) {
        return {
          inputTokens: result.data.usage.input_tokens ?? null,
          outputTokens: result.data.usage.output_tokens ?? null,
          cacheCreationInputTokens: result.data.usage.cache_creation_input_tokens ?? null,
          cacheReadInputTokens: result.data.usage.cache_read_input_tokens ?? null,
        };
      }
    } catch {
      // Fall through
    }
    return {
      inputTokens: null,
      outputTokens: null,
      cacheCreationInputTokens: null,
      cacheReadInputTokens: null,
    };
  },

  extractStream(
    raw: string,
    log: CapturedLog,
    fallbackModel?: string,
    collectChunks?: boolean,
  ): string {
    return extractAnthropicStream(raw, log, fallbackModel, collectChunks);
  },

  detectFormat(rawBody: string | null): boolean {
    if (rawBody === null) return false;
    try {
      const json: unknown = JSON.parse(rawBody);
      if (typeof json === "object" && json !== null && !Array.isArray(json)) {
        const keys = Object.keys(json);
        // Anthropic puts `system` as a top-level key alongside `model` and `messages`
        return keys.includes("model") && keys.includes("messages") && keys.includes("system");
      }
      return false;
    } catch {
      return false;
    }
  },
};
