import { createAmazonBedrock } from "@ai-sdk/amazon-bedrock";
import { createAnthropic } from "@ai-sdk/anthropic";
import { createCerebras } from "@ai-sdk/cerebras";
import { createCohere } from "@ai-sdk/cohere";
import { createGoogleGenerativeAI } from "@ai-sdk/google";
import { createVertex } from "@ai-sdk/google-vertex";
import { createGroq } from "@ai-sdk/groq";
import { createMistral } from "@ai-sdk/mistral";
import { createOpenAI } from "@ai-sdk/openai";
import { createXai } from "@ai-sdk/xai";
import type { InteractionSource } from "@archestra/shared";
import {
  CHAT_API_KEY_ID_HEADER,
  EXTERNAL_AGENT_ID_HEADER,
  PROVIDER_BASE_URL_HEADER,
  providerRequiresPerUserCredential,
  requiresOpenAiResponsesApi,
  SESSION_ID_HEADER,
  SOURCE_HEADER,
  type SupportedProvider,
  UNTRUSTED_CONTEXT_HEADER,
  USER_ID_HEADER,
} from "@archestra/shared";
import { context, propagation } from "@opentelemetry/api";
import type { streamText } from "ai";
import { isAnthropicNativeEndpoint } from "@/clients/anthropic-endpoint";
import { anthropicWorkloadIdentity } from "@/clients/anthropic-workload-identity";
import { isAzureOpenAiEntraIdEnabled } from "@/clients/azure-openai-credentials";
import {
  createAzureFetchWithApiVersion,
  normalizeAzureApiKey,
} from "@/clients/azure-url";
import {
  decodeBedrockSigV4Marker,
  getBedrockCredentialProvider,
  getBedrockRegion,
  isBedrockIamAuthEnabled,
} from "@/clients/bedrock-credentials";
import { isVertexAiEnabled } from "@/clients/gemini-client";
import { getLlmUpstreamDispatcher } from "@/clients/llm-upstream-dispatcher";
import { openRouterAttributionHeaders } from "@/clients/openrouter-attribution";
import { createResponseHealingFetch } from "@/clients/openrouter-response-healing";
import config from "@/config";
import logger from "@/logging";
import { ApiError } from "@/types";
import { resolveProviderApiKey } from "@/utils/llm-api-key-resolution";
import { LlmProviderAuthRequiredError } from "@/utils/llm-provider-auth-error";

/**
 * Placeholder API key for providers that don't require authentication (vLLM, Ollama).
 * The OpenAI SDK requires a non-empty apiKey string, so we pass this sentinel value.
 */
const KEYLESS_PROVIDER_API_KEY_PLACEHOLDER = "EMPTY";

/**
 * Note: vLLM and Ollama use the @ai-sdk/openai provider since they expose OpenAI-compatible APIs.
 * When creating a vLLM/Ollama model, we use createOpenAI with the respective base URL.
 */

/**
 * Type representing a model that can be passed to streamText/generateText
 */
export type LLMModel = Parameters<typeof streamText>[0]["model"];

/**
 * Check if API key is required for the given provider
 */
export function isApiKeyRequired(
  provider: SupportedProvider,
  apiKey: string | undefined,
): boolean {
  if (apiKey) return false;
  // Gemini with Vertex AI doesn't require an API key
  if (provider === "gemini" && isVertexAiEnabled()) return false;
  return !!providerModelConfigs[provider].apiKeyRequiredMessage;
}

/**
 * Create an LLM model that calls the provider API directly (not through LLM Proxy).
 * Use this for meta operations like title generation that don't need proxy features.
 */
export function createDirectLLMModel({
  provider,
  apiKey,
  modelName,
  baseUrl,
}: {
  provider: SupportedProvider;
  apiKey: string | undefined;
  modelName: string;
  baseUrl: string | null;
}): LLMModel {
  const cfg = providerModelConfigs[provider];
  if (!cfg) {
    throw new ApiError(400, `Unsupported provider: ${provider}`);
  }
  if (cfg.apiKeyRequiredMessage && !apiKey) {
    throw new ApiError(400, cfg.apiKeyRequiredMessage);
  }
  const resolvedBaseUrl = baseUrl ?? cfg.defaultBaseUrl;
  const baseURL =
    resolvedBaseUrl && cfg.proxiedPathSuffix
      ? `${resolvedBaseUrl}${cfg.proxiedPathSuffix}`
      : resolvedBaseUrl;
  return cfg.createModel({
    apiKey,
    modelName,
    baseURL,
    headers: providerHeaders(cfg),
    // Direct OpenRouter models bypass the proxy adapter, so heal the request
    // body here; the wrapper no-ops for non-healable requests.
    fetch: provider === "openrouter" ? createResponseHealingFetch() : undefined,
  });
}

/**
 * Create an LLM model for the specified provider, pointing to the LLM Proxy
 * Returns a model instance ready to use with streamText/generateText
 */
export function createLLMModel(params: {
  provider: SupportedProvider;
  apiKey: string | undefined;
  agentId: string;
  modelName: string;
  userId?: string;
  externalAgentId?: string;
  sessionId?: string;
  source?: InteractionSource;
  baseUrl: string | null;
  contextIsTrusted?: boolean;
  chatApiKeyId?: string;
}): LLMModel {
  const {
    provider,
    apiKey,
    agentId,
    modelName,
    userId,
    externalAgentId,
    sessionId,
    source,
    baseUrl,
    contextIsTrusted,
    chatApiKeyId,
  } = params;

  // Build headers for LLM Proxy
  const clientHeaders: Record<string, string> = {};
  if (externalAgentId) {
    clientHeaders[EXTERNAL_AGENT_ID_HEADER] = externalAgentId;
  }
  if (userId) {
    clientHeaders[USER_ID_HEADER] = userId;
  }
  if (sessionId) {
    clientHeaders[SESSION_ID_HEADER] = sessionId;
  }
  if (source) {
    clientHeaders[SOURCE_HEADER] = source;
  }
  // Only propagate the header when the caller has explicitly established that
  // context is unsafe. `undefined` means trust was not evaluated for this flow,
  // so we preserve the default trusted behavior.
  if (contextIsTrusted === false) {
    clientHeaders[UNTRUSTED_CONTEXT_HEADER] = "true";
  }
  if (baseUrl) {
    clientHeaders[PROVIDER_BASE_URL_HEADER] = baseUrl;
  }
  // Chat sends the raw provider secret to the proxy, so the proxy can't tie
  // the call to a chat_api_keys row. Forwarding the row ID here lets the
  // proxy look up per-key configuration (extraHeaders). Loopback-gated on
  // the proxy side; see CHAT_API_KEY_ID_HEADER.
  if (chatApiKeyId) {
    clientHeaders[CHAT_API_KEY_ID_HEADER] = chatApiKeyId;
    logger.info(
      { chatApiKeyId, provider },
      `[${provider}Proxy] chat attaching provider-api-key-id header`,
    );
  }

  const headers =
    Object.keys(clientHeaders).length > 0 ? clientHeaders : undefined;

  const cfg = providerModelConfigs[provider];
  const proxyBaseUrl = buildProxyBaseUrl(provider, agentId);
  const baseURL = cfg.proxiedPathSuffix
    ? `${proxyBaseUrl}${cfg.proxiedPathSuffix}`
    : proxyBaseUrl;

  return cfg.createModel({
    apiKey,
    modelName,
    baseURL,
    headers,
    fetch: createTracedFetch(),
  });
}

/**
 * Full helper to resolve API key and create LLM model.
 * Provider must be explicitly passed - callers can use detectProviderFromModel
 * as a fallback for backward compatibility with existing conversations.
 */
export async function createLLMModelForAgent(params: {
  organizationId: string;
  userId: string;
  agentId: string;
  model: string;
  provider: SupportedProvider;
  conversationId?: string | null;
  externalAgentId?: string;
  sessionId?: string;
  source?: InteractionSource;
  agentLlmApiKeyId?: string | null;
  contextIsTrusted?: boolean;
}): Promise<{
  model: LLMModel;
  provider: SupportedProvider;
  apiKeySource: string;
  /**
   * True when this resolves to genuine Anthropic (vs an Anthropic-compatible
   * endpoint behind a custom base URL serving a non-Claude model). Gates
   * Anthropic-only request-body features in chat normalization, mirroring the
   * proxy's `anthropic-beta` header gating so the two can't drift.
   */
  anthropicNativeEndpoint: boolean;
}> {
  const {
    organizationId,
    userId,
    agentId,
    model: modelName,
    provider,
    conversationId,
    externalAgentId,
    sessionId,
    source,
    agentLlmApiKeyId,
    contextIsTrusted,
  } = params;

  const {
    apiKey,
    source: apiKeySource,
    baseUrl,
    chatApiKeyId,
  } = await resolveProviderApiKey({
    organizationId,
    userId,
    provider,
    conversationId,
    agentLlmApiKeyId,
  });

  // Check if Gemini with Vertex AI (doesn't require API key)
  const isGeminiWithVertexAi = provider === "gemini" && isVertexAiEnabled();
  // Check if Bedrock with IAM auth (doesn't require API key)
  const isBedrockWithIamAuth =
    provider === "bedrock" && isBedrockIamAuthEnabled();
  // vLLM and Ollama typically don't require API keys
  const isVllm = provider === "vllm";
  const isOllama = provider === "ollama";
  const isAzureWithEntra =
    provider === "azure" && isAzureOpenAiEntraIdEnabled();
  const isAnthropicWithWif =
    provider === "anthropic" && anthropicWorkloadIdentity.isEnabled();

  logger.info(
    {
      apiKeySource,
      provider,
      isGeminiWithVertexAi,
      isBedrockWithIamAuth,
      isVllm,
      isOllama,
      isAzureWithEntra,
      isAnthropicWithWif,
    },
    "Using LLM provider API key",
  );

  if (
    !apiKey &&
    !isGeminiWithVertexAi &&
    !isBedrockWithIamAuth &&
    !isVllm &&
    !isOllama &&
    !isAzureWithEntra &&
    !isAnthropicWithWif
  ) {
    // Per-user providers (GitHub Copilot) need the acting user's own linked
    // account; surface a typed error so callers can prompt them to connect
    // rather than showing a generic "configure a key" message.
    if (providerRequiresPerUserCredential(provider)) {
      throw new LlmProviderAuthRequiredError(provider);
    }
    throw new ApiError(
      400,
      "LLM Provider API key not configured. Please configure it in Provider Settings.",
    );
  }

  const model = createLLMModel({
    provider,
    apiKey,
    agentId,
    modelName,
    userId,
    externalAgentId,
    sessionId,
    source,
    baseUrl,
    contextIsTrusted,
    chatApiKeyId,
  });

  const anthropicNativeEndpoint = isAnthropicNativeEndpoint({
    provider,
    model: modelName,
    baseUrl,
  });

  return { model, provider, apiKeySource, anthropicNativeEndpoint };
}

// =============================================================================
// Internal helpers
// =============================================================================

/**
 * Unified model creation config for each provider.
 * A single `createModel` function handles both direct and proxied calls.
 *
 * For direct calls: only apiKey, modelName, baseURL are provided.
 * For proxied calls: headers and fetch are also provided (for trace context injection).
 */
type ProviderModelConfig = {
  createModel: (params: {
    apiKey: string | undefined;
    modelName: string;
    baseURL: string | undefined;
    headers?: Record<string, string>;
    fetch?: typeof globalThis.fetch;
  }) => LLMModel;
  /** Default base URL for direct calls (falls back to provider's built-in default when undefined) */
  defaultBaseUrl: string | undefined;
  /** Error message when API key is missing. Undefined = key is optional (vllm, ollama). */
  apiKeyRequiredMessage?: string;
  /** Path suffix appended to proxy base URL for proxied calls (e.g. "/v1" for anthropic) */
  proxiedPathSuffix?: string;
  /** Static headers always sent to the provider (e.g. OpenRouter attribution). */
  extraHeaders?: Record<string, string>;
};

/** Static provider headers (e.g. OpenRouter attribution), or undefined when none. */
function providerHeaders(
  cfg: ProviderModelConfig,
): Record<string, string> | undefined {
  return cfg.extraHeaders && Object.keys(cfg.extraHeaders).length > 0
    ? cfg.extraHeaders
    : undefined;
}

/**
 * Unified registry of model configs for each provider.
 * TypeScript enforces that ALL providers in SupportedProvider have an entry.
 * Adding a new provider to SupportedProvider will cause a compile error here
 * until the corresponding config is added.
 */
const providerModelConfigs: Record<SupportedProvider, ProviderModelConfig> = {
  // --- Native SDK providers (use their own SDK, call client(modelName)) ---

  anthropic: {
    createModel: ({ apiKey, modelName, baseURL, headers, fetch }) =>
      createAnthropic({ apiKey, baseURL, headers, fetch })(modelName),
    defaultBaseUrl: config.llm.anthropic.baseUrl,
    apiKeyRequiredMessage:
      "Anthropic API key is required. Please configure ANTHROPIC_API_KEY.",
    proxiedPathSuffix: "/v1",
  },

  cerebras: {
    createModel: ({ apiKey, modelName, baseURL, headers, fetch }) =>
      createCerebras({ apiKey, baseURL, headers, fetch })(modelName),
    defaultBaseUrl: config.llm.cerebras.baseUrl,
    apiKeyRequiredMessage:
      "Cerebras API key is required. Please configure CEREBRAS_API_KEY.",
  },

  cohere: {
    createModel: ({ apiKey, modelName, baseURL, headers, fetch }) =>
      createCohere({ apiKey, baseURL, headers, fetch })(modelName),
    defaultBaseUrl: config.llm.cohere.baseUrl,
    apiKeyRequiredMessage:
      "Cohere API key is required. Please configure COHERE_API_KEY.",
  },

  mistral: {
    createModel: ({ apiKey, modelName, baseURL, headers, fetch }) =>
      createMistral({ apiKey, baseURL, headers, fetch })(modelName),
    defaultBaseUrl: config.llm.mistral.baseUrl,
    apiKeyRequiredMessage:
      "Mistral API key is required. Please configure MISTRAL_API_KEY.",
  },

  groq: {
    createModel: ({ apiKey, modelName, baseURL, headers, fetch }) =>
      createGroq({ apiKey, baseURL, headers, fetch })(modelName),
    defaultBaseUrl: config.llm.groq.baseUrl,
    apiKeyRequiredMessage:
      "Groq API key is required. Please configure ARCHESTRA_CHAT_GROQ_API_KEY.",
  },

  xai: {
    createModel: ({ apiKey, modelName, baseURL, headers, fetch }) =>
      createXai({ apiKey, baseURL, headers, fetch })(modelName),
    defaultBaseUrl: config.llm.xai.baseUrl,
    apiKeyRequiredMessage:
      "xAI API key is required. Please configure ARCHESTRA_CHAT_XAI_API_KEY.",
  },

  // --- OpenAI-compatible providers (use createOpenAI with .chat()) ---

  openai: {
    createModel: ({ apiKey, modelName, baseURL, headers, fetch }) => {
      const client = createOpenAI({ apiKey, baseURL, headers, fetch });
      // "pro" reasoning models are Responses-API-only; routing them through
      // .chat() hits /chat/completions and 404s. See requiresOpenAiResponsesApi.
      return requiresOpenAiResponsesApi(modelName)
        ? client.responses(modelName)
        : client.chat(modelName);
    },
    defaultBaseUrl: config.llm.openai.baseUrl,
    apiKeyRequiredMessage:
      "OpenAI API key is required. Please configure OPENAI_API_KEY.",
  },

  openrouter: {
    createModel: ({ apiKey, modelName, baseURL, headers, fetch }) =>
      createOpenAI({ apiKey, baseURL, headers, fetch }).chat(modelName),
    defaultBaseUrl: config.llm.openrouter.baseUrl,
    apiKeyRequiredMessage:
      "OpenRouter API key is required. Please configure ARCHESTRA_CHAT_OPENROUTER_API_KEY.",
    extraHeaders: openRouterAttributionHeaders(),
  },

  perplexity: {
    createModel: ({ apiKey, modelName, baseURL, headers, fetch }) =>
      createOpenAI({ apiKey, baseURL, headers, fetch }).chat(modelName),
    defaultBaseUrl: config.llm.perplexity.baseUrl,
    apiKeyRequiredMessage:
      "Perplexity API key is required. Please configure PERPLEXITY_API_KEY.",
  },

  zhipuai: {
    createModel: ({ apiKey, modelName, baseURL, headers, fetch }) =>
      createOpenAI({ apiKey, baseURL, headers, fetch }).chat(modelName),
    defaultBaseUrl: config.llm.zhipuai.baseUrl,
    apiKeyRequiredMessage:
      "Zhipu AI API key is required. Please configure ZHIPUAI_API_KEY.",
  },

  minimax: {
    createModel: ({ apiKey, modelName, baseURL, headers, fetch }) =>
      createOpenAI({ apiKey, baseURL, headers, fetch }).chat(modelName),
    defaultBaseUrl: config.llm.minimax.baseUrl,
    apiKeyRequiredMessage:
      "MiniMax API key is required. Please configure ARCHESTRA_CHAT_MINIMAX_API_KEY.",
  },

  deepseek: {
    createModel: ({ apiKey, modelName, baseURL, headers, fetch }) =>
      createOpenAI({ apiKey, baseURL, headers, fetch }).chat(modelName),
    defaultBaseUrl: config.llm.deepseek.baseUrl,
    apiKeyRequiredMessage:
      "DeepSeek API key is required. Please configure DEEPSEEK_API_KEY.",
  },

  "github-copilot": {
    // The model always talks to the local LLM proxy (buildProxyBaseUrl), and
    // the proxy's github-copilot adapter exchanges the GitHub OAuth token for
    // the short-lived Copilot bearer — exchanging here too would hand the
    // proxy an already-exchanged bearer it cannot exchange again.
    createModel: ({ apiKey, modelName, baseURL, headers, fetch }) =>
      createOpenAI({ apiKey, baseURL, headers, fetch }).chat(modelName),
    defaultBaseUrl: config.llm["github-copilot"].baseUrl,
    apiKeyRequiredMessage:
      "GitHub Copilot requires a GitHub OAuth token. Connect your GitHub account or configure ARCHESTRA_CHAT_GITHUB_COPILOT_API_KEY.",
  },

  azure: {
    createModel: ({
      apiKey,
      modelName,
      baseURL,
      headers,
      fetch: providedFetch,
    }) => {
      // The AI SDK client can't set Azure's api-version as a default query param,
      // so we wrap fetch and inject it on every request.
      const fetchWithVersion = createAzureFetchWithApiVersion({
        apiVersion: config.llm.azure.apiVersion,
        fetch: providedFetch,
      });
      const normalizedApiKey = normalizeAzureApiKey(apiKey);
      const sdkApiKey =
        normalizedApiKey ??
        (isAzureOpenAiEntraIdEnabled()
          ? KEYLESS_PROVIDER_API_KEY_PLACEHOLDER
          : undefined);
      return createOpenAI({
        apiKey: sdkApiKey,
        baseURL,
        headers: normalizedApiKey
          ? { ...headers, "api-key": normalizedApiKey }
          : headers,
        fetch: fetchWithVersion,
      }).chat(modelName);
    },
    defaultBaseUrl: config.llm.azure.baseUrl || undefined,
    apiKeyRequiredMessage:
      "Azure AI Foundry API key is required. Please configure ARCHESTRA_CHAT_AZURE_OPENAI_API_KEY.",
  },

  // --- OpenAI-compatible providers with optional API key ---

  vllm: {
    createModel: ({ apiKey, modelName, baseURL, headers, fetch }) =>
      createOpenAI({
        apiKey: apiKey || KEYLESS_PROVIDER_API_KEY_PLACEHOLDER,
        baseURL,
        headers,
        fetch,
      }).chat(modelName),
    defaultBaseUrl: config.llm.vllm.baseUrl,
    // No apiKeyRequiredMessage — key is optional
  },

  ollama: {
    createModel: ({ apiKey, modelName, baseURL, headers, fetch }) =>
      createOpenAI({
        apiKey: apiKey || KEYLESS_PROVIDER_API_KEY_PLACEHOLDER,
        baseURL,
        headers,
        fetch,
      }).chat(modelName),
    defaultBaseUrl: config.llm.ollama.baseUrl,
    // No apiKeyRequiredMessage — key is optional
  },

  // --- Special providers ---

  gemini: {
    createModel: ({ apiKey, modelName, baseURL, headers, fetch }) => {
      // Proxied path (headers/fetch provided): always use GoogleGenerativeAI
      if (headers || fetch) {
        return createGoogleGenerativeAI({
          apiKey: apiKey || "vertex-ai-mode",
          baseURL,
          headers,
          fetch,
        })(modelName);
      }
      // Direct path: use Vertex AI if enabled
      if (isVertexAiEnabled()) {
        const { vertexAi } = config.llm.gemini;
        return createVertex({
          project: vertexAi.project,
          location: vertexAi.location,
          googleAuthOptions: {
            projectId: vertexAi.project,
            ...(vertexAi.credentialsFile && {
              keyFilename: vertexAi.credentialsFile,
            }),
          },
        })(modelName);
      }
      // Direct path without Vertex AI — key is required
      if (!apiKey) {
        throw new ApiError(
          400,
          "Gemini API key is required when Vertex AI is not enabled. Please configure GEMINI_API_KEY or enable Vertex AI.",
        );
      }
      return createGoogleGenerativeAI({ apiKey, baseURL })(modelName);
    },
    defaultBaseUrl: undefined, // GoogleGenerativeAI has its own default
    // apiKeyRequiredMessage is undefined — validation is inside createModel (Vertex AI special case)
    proxiedPathSuffix: "/v1beta",
  },

  bedrock: {
    createModel: ({ apiKey, modelName, baseURL, headers, fetch }) => {
      const region = getBedrockRegion(baseURL);

      if (!apiKey && isBedrockIamAuthEnabled()) {
        return createAmazonBedrock({
          region,
          baseURL,
          credentialProvider: getBedrockCredentialProvider(),
          headers,
          fetch,
        })(modelName);
      }

      const sigV4 = decodeBedrockSigV4Marker(apiKey);
      if (sigV4) {
        return createAmazonBedrock({
          region,
          baseURL,
          accessKeyId: sigV4.accessKeyId,
          secretAccessKey: sigV4.secretAccessKey,
          sessionToken: sigV4.sessionToken,
          headers,
          fetch,
        })(modelName);
      }

      return createAmazonBedrock({
        apiKey,
        region,
        baseURL,
        secretAccessKey: undefined,
        accessKeyId: undefined,
        sessionToken: undefined,
        credentialProvider: undefined,
        headers,
        fetch,
      })(modelName);
    },
    defaultBaseUrl: config.llm.bedrock.baseUrl,
    apiKeyRequiredMessage: isBedrockIamAuthEnabled()
      ? undefined
      : "Amazon Bedrock API key is required. Please configure ARCHESTRA_CHAT_BEDROCK_API_KEY.",
  },
};

/**
 * Creates a fetch wrapper that injects W3C trace context (traceparent/tracestate)
 * into outgoing HTTP headers. This enables the LLM proxy handler to extract the
 * parent context and create child spans, linking chat → LLM proxy traces together.
 */
function createTracedFetch(): typeof globalThis.fetch {
  return (input, init) => {
    const headers = new Headers(init?.headers);
    // Inject active trace context (traceparent, tracestate) into outgoing headers.
    // Uses a carrier object because propagation.inject expects a plain object,
    // then copies the injected headers into the actual Headers instance.
    const carrier: Record<string, string> = {};
    propagation.inject(context.active(), carrier);
    for (const [key, value] of Object.entries(carrier)) {
      headers.set(key, value);
    }
    // Opt-in upstream timeout dispatcher; undefined leaves undici defaults
    // untouched. See @/clients/llm-upstream-dispatcher.
    const dispatcher = getLlmUpstreamDispatcher();

    if (!dispatcher) {
      return globalThis.fetch(input, { ...init, headers });
    }

    return globalThis.fetch(input, {
      ...init,
      headers,
      dispatcher,
    } as RequestInit);
  };
}

/**
 * Build the proxy base URL for a provider
 */
function buildProxyBaseUrl(provider: string, agentId: string): string {
  return `http://localhost:${config.api.port}/v1/${provider}/${agentId}`;
}
