// SPDX-FileCopyrightText: Amolith <amolith@secluded.site>
//
// SPDX-License-Identifier: GPL-3.0-or-later

import { getModel, type Model } from "@mariozechner/pi-ai";
import {
  type CustomModelConfig,
  type RumiloConfig,
} from "../config/schema.js";
import { ConfigError } from "../util/errors.js";
import { resolveHeaders } from "../util/env.js";

export function resolveModel(
  modelString: string,
  config: RumiloConfig,
): Model<any> {
  const colonIndex = modelString.indexOf(":");
  if (colonIndex === -1) {
    throw new ConfigError(`Invalid model format: "${modelString}". Expected provider:model (e.g. anthropic:claude-sonnet-4-20250514)`);
  }

  const provider = modelString.slice(0, colonIndex);
  const modelName = modelString.slice(colonIndex + 1);

  if (!provider || !modelName) {
    throw new ConfigError(`Invalid model format: "${modelString}". Expected provider:model (e.g. anthropic:claude-sonnet-4-20250514)`);
  }

  // Handle custom models
  if (provider === "custom") {
    return resolveCustomModel(modelName, config);
  }

  // Handle built-in providers
  return getModel(provider as any, modelName);
}

function resolveCustomModel(modelName: string, config: RumiloConfig): Model<any> {
  if (!config.custom_models) {
    throw new ConfigError(
      `No custom models defined in config. Add a [custom_models.${modelName}] section to use custom:${modelName}`,
    );
  }

  const customConfig = config.custom_models[modelName];
  if (!customConfig) {
    const available = Object.keys(config.custom_models).join(", ");
    throw new ConfigError(
      `Custom model '${modelName}' not found. Available custom models: ${available}`,
    );
  }

  return buildCustomModel(customConfig);
}

function buildCustomModel(config: CustomModelConfig): Model<any> {
  const api = config.api as any;

  const cost: any = {
    input: config.cost.input,
    output: config.cost.output,
  };

  if (config.cost.cache_read !== undefined) {
    cost.cacheRead = config.cost.cache_read;
  }

  if (config.cost.cache_write !== undefined) {
    cost.cacheWrite = config.cost.cache_write;
  }

  const model: any = {
    id: config.id,
    name: config.name,
    api,
    provider: config.provider as any,
    baseUrl: config.base_url,
    reasoning: config.reasoning,
    input: config.input,
    cost,
    contextWindow: config.context_window,
    maxTokens: config.max_tokens,
  };

  const resolvedHeaders = resolveHeaders(config.headers);
  if (resolvedHeaders) {
    model.headers = resolvedHeaders;
  }

  if (config.compat) {
    model.compat = convertCompatConfig(config.compat);
  }

  return model;
}

function convertCompatConfig(
  compat: CustomModelConfig["compat"],
): any {
  if (!compat) {
    throw new Error("Compat config is expected to be defined");
  }

  return {
    supportsStore: compat.supports_store,
    supportsDeveloperRole: compat.supports_developer_role,
    supportsReasoningEffort: compat.supports_reasoning_effort,
    supportsUsageInStreaming: compat.supports_usage_in_streaming,
    maxTokensField: compat.max_tokens_field,
    requiresToolResultName: compat.requires_tool_result_name,
    requiresAssistantAfterToolResult: compat.requires_assistant_after_tool_result,
    requiresThinkingAsText: compat.requires_thinking_as_text,
    requiresMistralToolIds: compat.requires_mistral_tool_ids,
    thinkingFormat: compat.thinking_format,
  };
}