import {
  type Api,
  type AssistantMessage,
  type AssistantMessageEvent,
  type Context,
  type Credential,
  createAssistantMessageEventStream,
  type Model,
  type Provider,
  type ProviderHeaders,
  type RefreshModelsContext,
  type SimpleStreamOptions,
  type StreamOptions,
} from "@earendil-works/pi-ai";

function withProvider<T extends { provider: string }>(value: T, provider: string): T {
  return { ...value, provider };
}

function mapEvent(event: AssistantMessageEvent, provider: string): AssistantMessageEvent {
  switch (event.type) {
    case "start":
    case "text_start":
    case "text_delta":
    case "text_end":
    case "thinking_start":
    case "thinking_delta":
    case "thinking_end":
    case "toolcall_start":
    case "toolcall_delta":
    case "toolcall_end":
      return { ...event, partial: withProvider(event.partial, provider) };
    case "done":
      return { ...event, message: withProvider(event.message, provider) };
    case "error":
      return { ...event, error: withProvider(event.error, provider) };
  }
}

function remapStream(
  sourceStream: ReturnType<Provider["stream"]> | Promise<ReturnType<Provider["stream"]>>,
  aliasModel: Model<Api>,
  aliasId: string,
) {
  const stream = createAssistantMessageEventStream();

  void (async () => {
    try {
      for await (const event of await sourceStream) {
        stream.push(mapEvent(event, aliasId));
      }
      stream.end();
    } catch (error) {
      const message: AssistantMessage = {
        role: "assistant",
        content: [],
        api: aliasModel.api,
        provider: aliasId,
        model: aliasModel.id,
        usage: {
          input: 0,
          output: 0,
          cacheRead: 0,
          cacheWrite: 0,
          totalTokens: 0,
          cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
        },
        stopReason: "error",
        errorMessage: error instanceof Error ? error.message : String(error),
        timestamp: Date.now(),
      };
      stream.push({ type: "error", reason: "error", error: message });
      stream.end();
    }
  })();

  return stream;
}

function toSourceModel(model: Model<Api>, sourceId: string): Model<Api> {
  return withProvider(model, sourceId);
}

function toAliasModels(models: readonly Model<Api>[], aliasId: string): Model<Api>[] {
  return models.map((model) => withProvider(model, aliasId));
}

function mergeHeaders(
  base: ProviderHeaders | undefined,
  override: ProviderHeaders | undefined,
): ProviderHeaders | undefined {
  if (!base && !override) {
    return undefined;
  }

  const merged: ProviderHeaders = { ...base };
  for (const [name, value] of Object.entries(override ?? {})) {
    const lowerName = name.toLowerCase();
    for (const existingName of Object.keys(merged)) {
      if (existingName.toLowerCase() === lowerName) {
        delete merged[existingName];
      }
    }
    merged[name] = value;
  }
  return merged;
}

function withSourceHeaders<TOptions>(
  options: TOptions | undefined,
  sourceHeaders: ProviderHeaders | undefined,
): TOptions | undefined {
  if (!sourceHeaders) {
    return options;
  }

  return {
    ...options,
    headers: mergeHeaders(
      sourceHeaders,
      (options as { headers?: ProviderHeaders } | undefined)?.headers,
    ),
  } as TOptions;
}

function toSourceContext(context: Context, aliasId: string, sourceId: string): Context {
  return {
    ...context,
    messages: context.messages.map((message) =>
      message.role === "assistant" && message.provider === aliasId
        ? withProvider(message, sourceId)
        : message,
    ),
  };
}

export type SourceHeadersResolver = (model: Model<Api>) => Promise<ProviderHeaders | undefined>;

function forwardStream<TOptions>(
  resolveSource: () => Provider,
  aliasModel: Model<Api>,
  context: Context,
  options: TOptions | undefined,
  aliasId: string,
  invoke: (
    source: Provider,
    sourceModel: Model<Api>,
    sourceContext: Context,
    options: TOptions | undefined,
  ) => ReturnType<Provider["stream"]>,
  resolveSourceHeaders: SourceHeadersResolver | undefined,
) {
  const source = resolveSource();
  const sourceModel = toSourceModel(aliasModel, source.id);
  const sourceContext = toSourceContext(context, aliasId, source.id);
  const invokeSource = (sourceHeaders: ProviderHeaders | undefined) =>
    invoke(source, sourceModel, sourceContext, withSourceHeaders(options, sourceHeaders));
  const sourceStream = resolveSourceHeaders
    ? Promise.resolve()
        .then(() => resolveSourceHeaders(sourceModel))
        .then(invokeSource)
    : invokeSource(undefined);
  return remapStream(sourceStream, aliasModel, aliasId);
}

/**
 * Create a live facade over a provider while keeping the alias identity visible.
 *
 * Dynamic sources receive the alias-scoped refresh context, including its
 * credential and model store. The source provider still owns the catalog state,
 * because the Provider API does not expose a cloning or stateless refresh hook.
 */
export function createProviderAlias(
  sourceOrResolver: Provider | (() => Provider | undefined),
  aliasId: string,
  displayName?: string,
  resolveSourceHeaders?: SourceHeadersResolver,
): Provider {
  const resolveSource =
    typeof sourceOrResolver === "function" ? sourceOrResolver : () => sourceOrResolver;
  const initialSource = resolveSource();
  if (!initialSource) {
    throw new Error(`Provider "${aliasId}" source is unavailable.`);
  }

  const resolveRequiredSource = (): Provider => {
    const source = resolveSource();
    if (!source) {
      throw new Error(`Provider "${aliasId}" source is unavailable.`);
    }
    return source;
  };
  const stream = (model: Model<Api>, context: Context, options?: StreamOptions) =>
    forwardStream(
      resolveRequiredSource,
      model,
      context,
      options,
      aliasId,
      (source, sourceModel, sourceContext, sourceOptions) =>
        source.stream(sourceModel, sourceContext, sourceOptions),
      resolveSourceHeaders,
    );
  const streamSimple = (model: Model<Api>, context: Context, options?: SimpleStreamOptions) =>
    forwardStream(
      resolveRequiredSource,
      model,
      context,
      options,
      aliasId,
      (source, sourceModel, sourceContext, sourceOptions) =>
        source.streamSimple(sourceModel, sourceContext, sourceOptions),
      resolveSourceHeaders,
    );
  let refreshedSource: Provider | undefined;
  let refreshedModels: Model<Api>[] | undefined;
  const getModels = () => {
    const source = resolveRequiredSource();
    return refreshedSource === source && refreshedModels
      ? refreshedModels
      : toAliasModels(source.getModels(), aliasId);
  };
  const refreshModels = async (context: RefreshModelsContext): Promise<void> => {
    const source = resolveRequiredSource();
    if (!source.refreshModels) {
      return;
    }
    await source.refreshModels(context);
    refreshedSource = source;
    refreshedModels = toAliasModels(source.getModels(), aliasId);
  };

  return {
    id: aliasId,
    name: displayName ?? `${initialSource.name} (${aliasId})`,
    get baseUrl() {
      return resolveRequiredSource().baseUrl;
    },
    get headers() {
      return resolveRequiredSource().headers;
    },
    get auth() {
      return resolveRequiredSource().auth;
    },
    getModels,
    get refreshModels() {
      return resolveRequiredSource().refreshModels ? refreshModels : undefined;
    },
    filterModels: (models, credential: Credential | undefined) => {
      const source = resolveRequiredSource();
      return source.filterModels
        ? toAliasModels(
            source.filterModels(
              models.map((model) => toSourceModel(model, source.id)),
              credential,
            ),
            aliasId,
          )
        : models;
    },
    stream,
    streamSimple,
  };
}