/**
 * LLM Proxy Common Types
 *
 * These types define adapter interfaces for working with provider-specific
 * requests and responses in a uniform way. The original provider data is
 * preserved and can be reconstructed after modifications.
 *
 * Usage flow:
 * ```
 * Provider Request
 *       ↓
 * RequestAdapter (wraps original, provides uniform read/modify API)
 *       ↓
 * [Business Logic operates via adapter methods]
 *       ↓
 * adapter.toProviderRequest() → Modified Provider Request
 *       ↓
 * LLM Provider
 *       ↓
 * Provider Response
 *       ↓
 * ResponseAdapter (wraps original, provides uniform read API)
 *       ↓
 * [Business Logic operates via adapter methods]
 *       ↓
 * adapter.toProviderResponse() or adapter.toRefusalResponse()
 * ```

 */

import type {
  ArchestraInternalErrorCode,
  InteractionSource,
  SupportedProvider,
  SupportedProviderDiscriminator,
} from "@archestra/shared";

import type { Agent } from "./agent";

/**
 * GenAI operation names for tracing span names.
 * Follows OTEL GenAI Semantic Conventions.
 * Span names are constructed as `{operationName} {model}`.
 */
export type GenAiOperationName = "chat" | "generate_content" | "embedding";

import type {
  CommonMcpToolDefinition,
  CommonMessage,
  CommonToolCall,
  CommonToolResult,
} from "./common-llm-format";
import type { ToolCompressionStats } from "./tool-result-compression";

/**
 * Options for creating an LLM provider client
 */
export interface CreateClientOptions {
  /** Base URL override for the provider API */
  baseUrl?: string;
  /** Agent for observability metrics (request duration, tokens) */
  agent?: Agent;
  /** Default headers to include with every request */
  defaultHeaders?: Record<string, string>;
  /** Interaction source for observability metrics (e.g. "api", "chat", "knowledge:embedding") */
  source: InteractionSource;
}

/**
 * Adapter interface for LLM requests
 *
 * Wraps provider-specific request and provides uniform API for business logic.
 * Original data is preserved and can be reconstructed after modifications.
 *
 * @typeParam TRequest - Provider-specific request type
 * @typeParam TMessages - Provider-specific messages type
 */
export interface LLMRequestAdapter<TRequest, TMessages = unknown> {
  /** Provider name */
  readonly provider: SupportedProvider;

  // ---------------------------------------------------------------------------
  // Read Access
  // ---------------------------------------------------------------------------

  /** Get model name */
  getModel(): string;

  /** Check if streaming is requested */
  isStreaming(): boolean;

  /** Get messages in common format (for trusted data evaluation) */
  getMessages(): CommonMessage[];

  /** Get tool results from messages (for trusted data evaluation) */
  getToolResults(): CommonToolResult[];

  /** Get tool definitions (for persistence, hasTools check) */
  getTools(): CommonMcpToolDefinition[];

  /** Check if request has tools */
  hasTools(): boolean;

  /** Get provider-specific messages (for token counting) */
  getProviderMessages(): TMessages;

  /** Get original unmodified request */
  getOriginalRequest(): TRequest;

  // ---------------------------------------------------------------------------
  // Modify Access
  // ---------------------------------------------------------------------------

  /** Set model (for cost optimization) */
  setModel(model: string): void;

  /**
   * Update a tool result's content (for trusted data updates, TOON compression)
   * @param toolCallId - The tool call ID to update
   * @param newContent - New content string
   */
  updateToolResult(toolCallId: string, newContent: string): void;

  /**
   * Apply multiple tool result updates at once
   * @param updates - Map of tool call ID to new content
   */
  applyToolResultUpdates(updates: Record<string, string>): void;

  /**
   * Apply TOON compression to tool results
   * @param model - Model name for token counting
   * @returns Compression statistics
   *
   * TODO: Refactor to remove TOON logic from adapter. Instead:
   * 1. Calculate TOON updates externally: calculateToonUpdates(adapter.getToolResults(), model) → { updates, stats }
   * 2. Apply via existing applyToolResultUpdates(updates)
   * This keeps adapters simple (just apply updates) and makes TOON logic provider-agnostic.
   */
  applyToonCompression(model: string): Promise<ToolCompressionStats>;

  /**
   * Convert tool result content to provider-specific format (e.g., MCP image blocks)
   * @param messages - Provider-specific messages to convert
   */
  convertToolResultContent(messages: TMessages): TMessages;

  // ---------------------------------------------------------------------------
  // Build Modified Request
  // ---------------------------------------------------------------------------

  /**
   * Build the modified provider-specific request
   * Incorporates all modifications (model, tool results)
   */
  toProviderRequest(): TRequest;
}

// =============================================================================
// RESPONSE ADAPTER INTERFACE
// =============================================================================

/**
 * Adapter interface for LLM responses
 *
 * Wraps provider-specific response and provides uniform API for business logic.
 *
 * @typeParam TResponse - Provider-specific response type
 */
export interface LLMResponseAdapter<TResponse> {
  /** Provider name */
  readonly provider: SupportedProvider;

  // ---------------------------------------------------------------------------
  // Read Access
  // ---------------------------------------------------------------------------

  /** Get response ID */
  getId(): string;

  /** Get model name */
  getModel(): string;

  /** Get text content from response */
  getText(): string;

  /** Get tool calls from response (for tool invocation policies) */
  getToolCalls(): CommonToolCall[];

  /** Check if response has tool calls */
  hasToolCalls(): boolean;

  /** Get token usage */
  getUsage(): UsageView;

  /** Get original response */
  getOriginalResponse(): TResponse;

  /**
   * Optional hook for adapters whose wire response shape differs from the
   * shape we want to persist in the interaction log. When present, the proxy
   * handler stores this value instead of `getOriginalResponse()`. Default
   * (unimplemented) means log == wire, which is correct for every native
   * adapter.
   */
  getLoggedResponse?(): TResponse;

  /** Get finish reasons array for OTEL tracing (e.g., ["stop"], ["tool_calls"]) */
  getFinishReasons(): string[];

  // ---------------------------------------------------------------------------
  // Build Responses
  // ---------------------------------------------------------------------------

  /**
   * Build a refusal response (when tool invocation is blocked)
   * @param refusalMessage - Full message with metadata
   * @param contentMessage - Human-readable message
   */
  toRefusalResponse(refusalMessage: string, contentMessage: string): TResponse;
}

// =============================================================================
// STREAMING ADAPTER INTERFACE
// =============================================================================

/**
 * Accumulated state during streaming
 */
export interface StreamAccumulatorState {
  responseId: string;
  model: string;
  text: string;
  toolCalls: Array<{
    id: string;
    name: string;
    arguments: string;
  }>;
  /** Raw tool call events stored for replay after policy approval */
  rawToolCallEvents: unknown[];
  usage: UsageView | null;
  stopReason: string | null;
  // timing for metrics
  timing: {
    startTime: number;
    firstChunkTime: number | null;
  };
}

/**
 * Result of processing a stream chunk
 */
export interface ChunkProcessingResult {
  /** SSE data to send to client immediately (null if should be held) */
  sseData: string | Uint8Array | null;
  /** Whether this chunk contains tool call data (held for policy evaluation) */
  isToolCallChunk: boolean;
  /** Whether this is the final chunk */
  isFinal: boolean;
  /** Error information if this chunk represents an error event */
  error?: {
    type: string;
    message: string;
  };
}

/**
 * Adapter interface for streaming LLM responses
 *
 * Handles parsing provider-specific chunks, accumulating state,
 * and formatting SSE events.
 *
 * @typeParam TChunk - Provider-specific stream chunk type
 * @typeParam TResponse - Provider-specific response type
 */
export interface LLMStreamAdapter<TChunk, TResponse> {
  /** Provider name */
  readonly provider: SupportedProvider;

  /** Current accumulated state */
  readonly state: StreamAccumulatorState;

  // ---------------------------------------------------------------------------
  // Chunk Processing
  // ---------------------------------------------------------------------------

  /**
   * processChunk process straming chunkgs one-by-one to:
   * 1. Tell LLMProxy how to handle current streaming chunk, depending on it's type (stream immediately, buffer for tool result inspection, etc.)
   * 2. Updates concrete StreamAdapter's internal StreamingAccumulatorState. This state is also used by the LLMProxy
   *    to build final response and other business logic.
   */
  processChunk(chunk: TChunk): ChunkProcessingResult;

  // ---------------------------------------------------------------------------
  // SSE Formatting
  // ---------------------------------------------------------------------------

  /** Format SSE headers for response */
  getSSEHeaders(): Record<string, string>;

  /**
   * Format a text fragment as SSE to inject into an ongoing stream.
   * Used for progress messages (e.g., dual LLM status) during streaming.
   */
  formatTextDeltaSSE(text: string): string | Uint8Array;

  /** Get raw tool call events as SSE strings (for replay after policy approval) */
  getRawToolCallEvents(): (string | Uint8Array)[];

  /**
   * Format a complete, self-contained text response as SSE events.
   * Used when replacing the response entirely (e.g., policy refusal).
   * Returns provider-specific events that form a valid complete response.
   */
  formatCompleteTextSSE(text: string): (string | Uint8Array)[];

  /** Format the stream end marker */
  formatEndSSE(): string | Uint8Array;

  // ---------------------------------------------------------------------------
  // Build Response
  // ---------------------------------------------------------------------------

  /**
   * Reconstructs a complete provider-native response from accumulated streaming chunks.
   *
   * During streaming, responses arrive as many small chunks (text deltas, tool call fragments, etc.).
   * This method combines all accumulated state into a single complete response object,
   * which is needed for saving the interaction to the database.
   */
  toProviderResponse(): TResponse;
}

// =============================================================================
// ADAPTER FACTORY INTERFACE
// =============================================================================

/**
 * Factory for creating adapters for a specific provider
 *
 * Each provider implements this interface to create adapters for their
 * request/response types.
 *
 * @typeParam TRequest - Provider-specific request type
 * @typeParam TResponse - Provider-specific response type
 * @typeParam TMessages - Provider-specific messages type
 * @typeParam TChunk - Provider-specific stream chunk type
 * @typeParam THeaders - Provider-specific headers type
 */
export interface LLMProvider<TRequest, TResponse, TMessages, TChunk, THeaders> {
  /** Provider name */
  readonly provider: SupportedProvider;

  /** Interaction type for database storage */
  readonly interactionType: SupportedProviderDiscriminator;

  // ---------------------------------------------------------------------------
  // Adapter Creation
  // ---------------------------------------------------------------------------

  /** Create a request adapter */
  createRequestAdapter(
    request: TRequest,
  ): LLMRequestAdapter<TRequest, TMessages>;

  /** Create a response adapter */
  createResponseAdapter(response: TResponse): LLMResponseAdapter<TResponse>;

  /** Create a stream adapter. Request is optional and used by some providers (e.g., Bedrock for tool name mapping) */
  createStreamAdapter(request?: TRequest): LLMStreamAdapter<TChunk, TResponse>;

  // ---------------------------------------------------------------------------
  // Client & Headers
  // ---------------------------------------------------------------------------

  /** Extract API key from headers */
  extractApiKey(headers: THeaders): string | undefined;

  /** Get base URL for the provider (from config), undefined means use SDK default */
  getBaseUrl(): string | undefined;

  /** GenAI operation name for tracing (e.g., "chat", "generate_content"). The span name is constructed as `{operationName} {model}` by startActiveLlmSpan. */
  readonly spanName: GenAiOperationName;

  /**
   * Create provider client with observability.
   * Each provider is responsible for setting up its own metrics tracking:
   */
  createClient(
    apiKey: string | undefined,
    options: CreateClientOptions,
  ): unknown;

  // ---------------------------------------------------------------------------
  // Execution
  // ---------------------------------------------------------------------------

  /** Execute non-streaming request */
  execute(client: unknown, request: TRequest): Promise<TResponse>;

  /** Execute streaming request */
  executeStream(
    client: unknown,
    request: TRequest,
  ): Promise<AsyncIterable<TChunk>>;

  /**
   * Extract error message from provider-specific SDK error.
   * Each provider SDK wraps errors differently (e.g., Anthropic uses nested
   * error.error.message structure), so this normalizes them to a string.
   */
  extractErrorMessage(error: unknown): string;

  /**
   * Classify a provider-specific SDK error into an Archestra-normalized
   * internal code. Used by the proxy to emit a uniform `internal_code`
   * field on the error response body, which downstream consumers (notably
   * the chat-error mapper) can read without provider-specific knowledge.
   *
   * Returns `undefined` when the error doesn't match any known normalized
   * category, in which case no `internal_code` is added to the envelope.
   */
  extractInternalCode(error: unknown): ArchestraInternalErrorCode | undefined;
}

/**
 * Token usage from response
 */
export interface UsageView {
  /** Uncached input tokens only (normalized across providers). */
  inputTokens: number;
  outputTokens: number;
  /** Tokens served from the prompt cache (billed at the provider's reduced read rate). Absent = no cache reporting. */
  cacheReadTokens?: number;
  /** Tokens written to the prompt cache (0/absent for providers that auto-cache without a write surcharge). */
  cacheWriteTokens?: number;
  /** Portion of cacheWriteTokens written at the 1-hour TTL (billed higher than 5m). Anthropic-only; absent elsewhere. */
  cacheWrite1hTokens?: number;
  /** Output tokens spent on reasoning/extended thinking. Reported by OpenAI (reasoning_tokens) and Gemini (thoughtsTokenCount); absent elsewhere. */
  reasoningTokens?: number;
  /** True when inputTokens is a locally-computed estimate, not a provider-reported count (see applyInputTokenFallback). Absent/false = provider-measured. */
  inputTokensEstimated?: boolean;
}

/**
 * Create initial stream accumulator state
 */
export function createStreamAccumulatorState(): StreamAccumulatorState {
  return {
    responseId: "",
    model: "",
    text: "",
    toolCalls: [],
    rawToolCallEvents: [],
    usage: null,
    stopReason: null,
    timing: {
      startTime: Date.now(),
      firstChunkTime: null,
    },
  };
}
