import { OpenAICompatibleChatLanguageModel } from "@ai-sdk/openai-compatible";
import { OpenAIResponsesLanguageModel } from "@ai-sdk/openai/internal";
import {
  FetchFunction,
  loadApiKey,
  withoutTrailingSlash,
} from "@ai-sdk/provider-utils";

import log from "electron-log";
import { getExtraProviderOptions } from "./thinking_utils";
import { DYAD_INTERNAL_REQUEST_ID_HEADER } from "./provider_options";
import type { UserSettings } from "../../lib/schemas";
import type { LanguageModel } from "ai";

const logger = log.scope("llm_engine_provider");

export type ExampleChatModelId = string & {};
export interface ChatParams {
  providerId: string;
}
export interface ExampleProviderSettings {
  /**
Example API key.
*/
  apiKey?: string;
  /**
Base URL for the API calls.
*/
  baseURL?: string;
  /**
Custom headers to include in the requests.
*/
  headers?: Record<string, string>;
  /**
Optional custom url query parameters to include in request urls.
*/
  queryParams?: Record<string, string>;
  /**
Custom fetch implementation. You can use it as a middleware to intercept requests,
or to provide a custom fetch implementation for e.g. testing.
*/
  fetch?: FetchFunction;

  dyadOptions: {
    enableLazyEdits?: boolean;
    enableSmartFilesContext?: boolean;
    enableWebSearch?: boolean;
  };
  settings: UserSettings;
}

export interface DyadEngineProvider {
  /**
Creates a model for text generation.
*/
  (modelId: ExampleChatModelId, chatParams: ChatParams): LanguageModel;

  /**
Creates a chat model for text generation.
*/
  chatModel(modelId: ExampleChatModelId, chatParams: ChatParams): LanguageModel;

  responses(modelId: ExampleChatModelId, chatParams: ChatParams): LanguageModel;
}

export function createDyadEngine(
  options: ExampleProviderSettings,
): DyadEngineProvider {
  const baseURL = withoutTrailingSlash(options.baseURL);
  logger.info("creating dyad engine with baseURL", baseURL);

  // Track request ID attempts
  const requestIdAttempts = new Map<string, number>();

  const getHeaders = () => ({
    Authorization: `Bearer ${loadApiKey({
      apiKey: options.apiKey,
      environmentVariableName: "DYAD_PRO_API_KEY",
      description: "Example API key",
    })}`,
    ...options.headers,
  });

  interface CommonModelConfig {
    provider: string;
    url: ({ path }: { path: string }) => string;
    headers: () => Record<string, string>;
    fetch?: FetchFunction;
  }

  const getCommonModelConfig = (): CommonModelConfig => ({
    provider: `dyad-engine`,
    url: ({ path }) => {
      const url = new URL(`${baseURL}${path}`);
      if (options.queryParams) {
        url.search = new URLSearchParams(options.queryParams).toString();
      }
      return url.toString();
    },
    headers: getHeaders,
    fetch: options.fetch,
  });

  // Custom fetch implementation that adds dyad-specific options to the request
  const createDyadFetch = ({
    providerId,
  }: {
    providerId: string;
  }): FetchFunction => {
    return (input: RequestInfo | URL, init?: RequestInit) => {
      // Use default fetch if no init or body
      if (!init || !init.body || typeof init.body !== "string") {
        return (options.fetch || fetch)(input, init);
      }

      try {
        // Parse the request body to manipulate it
        const parsedBody = {
          ...JSON.parse(init.body),
          ...getExtraProviderOptions(providerId, options.settings),
        };

        const dyadVersionedFiles = parsedBody.dyadVersionedFiles;
        if ("dyadVersionedFiles" in parsedBody) {
          delete parsedBody.dyadVersionedFiles;
        }
        const dyadFiles = parsedBody.dyadFiles;
        if ("dyadFiles" in parsedBody) {
          delete parsedBody.dyadFiles;
        }
        // Read from body (OpenAICompatible models spread providerOptions into
        // the body) with a fallback to an internal header (OpenAIResponses
        // models don't forward providerOptions, so we pass it via header).
        const requestId =
          parsedBody.dyadRequestId ??
          (init.headers as Record<string, string> | undefined)?.[
            DYAD_INTERNAL_REQUEST_ID_HEADER
          ];
        if ("dyadRequestId" in parsedBody) {
          delete parsedBody.dyadRequestId;
        }
        const dyadAppId = parsedBody.dyadAppId;
        if ("dyadAppId" in parsedBody) {
          delete parsedBody.dyadAppId;
        }
        const dyadDisableFiles = parsedBody.dyadDisableFiles;
        if ("dyadDisableFiles" in parsedBody) {
          delete parsedBody.dyadDisableFiles;
        }
        const dyadMentionedApps = parsedBody.dyadMentionedApps;
        if ("dyadMentionedApps" in parsedBody) {
          delete parsedBody.dyadMentionedApps;
        }
        const dyadSmartContextMode = parsedBody.dyadSmartContextMode;
        if ("dyadSmartContextMode" in parsedBody) {
          delete parsedBody.dyadSmartContextMode;
        }

        // Track and modify requestId with attempt number
        let modifiedRequestId = requestId;
        if (requestId) {
          const currentAttempt = (requestIdAttempts.get(requestId) || 0) + 1;
          requestIdAttempts.set(requestId, currentAttempt);
          modifiedRequestId = `${requestId}:attempt-${currentAttempt}`;
        }

        // Add files to the request if they exist
        if (!dyadDisableFiles) {
          parsedBody.dyad_options = {
            files: dyadFiles,
            versioned_files: dyadVersionedFiles,
            enable_lazy_edits: options.dyadOptions.enableLazyEdits,
            enable_smart_files_context:
              options.dyadOptions.enableSmartFilesContext,
            smart_context_mode: dyadSmartContextMode,
            enable_web_search: options.dyadOptions.enableWebSearch,
            app_id: dyadAppId,
          };
          if (dyadMentionedApps?.length) {
            parsedBody.dyad_options.mentioned_apps = dyadMentionedApps;
          }
        }

        // Return modified request with files included and requestId in headers
        const { [DYAD_INTERNAL_REQUEST_ID_HEADER]: _, ...outgoingHeaders } =
          (init.headers as Record<string, string>) ?? {};
        const modifiedInit = {
          ...init,
          headers: {
            ...outgoingHeaders,
            ...(modifiedRequestId && {
              "X-Dyad-Request-Id": modifiedRequestId,
            }),
          },
          body: JSON.stringify(parsedBody),
        };

        // Use the provided fetch or default fetch
        return (options.fetch || fetch)(input, modifiedInit);
      } catch (e) {
        logger.error("Error parsing request body", e);
        // If parsing fails, use original request
        return (options.fetch || fetch)(input, init);
      }
    };
  };

  const createChatModel = (
    modelId: ExampleChatModelId,
    chatParams: ChatParams,
  ) => {
    const config = {
      ...getCommonModelConfig(),
      fetch: createDyadFetch({ providerId: chatParams.providerId }),
    };

    return new OpenAICompatibleChatLanguageModel(modelId, config);
  };

  const createResponsesModel = (
    modelId: ExampleChatModelId,
    chatParams: ChatParams,
  ) => {
    const config = {
      ...getCommonModelConfig(),
      fetch: createDyadFetch({ providerId: chatParams.providerId }),
    };

    return new OpenAIResponsesLanguageModel(modelId, config);
  };

  const provider = (modelId: ExampleChatModelId, chatParams: ChatParams) =>
    createChatModel(modelId, chatParams);

  provider.chatModel = createChatModel;
  provider.responses = createResponsesModel;

  return provider;
}
