gongxue-base/apps/server/src/ai-config/ai-config.service.ts

import {
  BadRequestException,
  Injectable,
  InternalServerErrorException,
  Logger,
} from '@nestjs/common';
import { InjectRepository } from '@nestjs/typeorm';
import { Repository } from 'typeorm';
import { AiConfig, AiProvider, SINGLETON_KEY } from './ai-config.entity';
import {
  AiConfigResponseDto,
  AiRuntimeConfig,
  DEFAULT_BASE_URLS,
  FetchModelsDto,
  FetchModelsResultDto,
  SaveAiConfigDto,
  TestAiConfigDto,
  AiConfigTestResultDto,
} from './dto/ai-config.dto';
import { decrypt, encrypt, validateAndNormalizeBaseUrl, validateDnsNotPrivate } from './ai-config.helpers';
import { fetchModels, testConnection } from './ai-config.probe';
import type { AiConfigProbeContext } from './ai-config.probe';

@Injectable()
export class AiConfigService implements AiConfigProbeContext {
  private readonly logger = new Logger(AiConfigService.name);

  constructor(
    @InjectRepository(AiConfig)
    private readonly repo: Repository<AiConfig>,
  ) {}

  save(config: AiConfig): Promise<AiConfig> {
    return this.repo.save(config);
  }

  /** Resolve the effective API key: DB first, then env, then none */
  resolveApiKey(config: AiConfig | null): {
    plaintext: string | null;
    source: 'database' | 'environment' | 'none';
  } {
    // DB stored key
    if (config?.encryptedApiKey && config?.apiKeyIv && config?.apiKeyAuthTag) {
      try {
        const plaintext = decrypt(config.encryptedApiKey, config.apiKeyIv, config.apiKeyAuthTag);
        return { plaintext, source: 'database' };
      } catch {
        this.logger.error('解密数据库 API Key 失败，密文可能已损坏');
        throw new InternalServerErrorException('无法解密 API Key');
      }
    }

    // Environment fallback
    const envKey = process.env.AI_API_KEY;
    if (envKey) {
      return { plaintext: envKey, source: 'environment' };
    }

    return { plaintext: null, source: 'none' };
  }

  /** Load or create the singleton config row */
  async getOrCreateConfig(): Promise<AiConfig> {
    let config = await this.repo.findOne({ where: { singletonKey: SINGLETON_KEY } });
    if (!config) {
      config = this.repo.create({
        singletonKey: SINGLETON_KEY,
        provider: AiProvider.DEEPSEEK,
        baseUrl: DEFAULT_BASE_URLS[AiProvider.DEEPSEEK],
        enabled: true,
        timeoutMs: 30000,
      });
      try {
        config = await this.repo.save(config);
      } catch (err: unknown) {
        // Unique constraint violation → another request created it first
        const isErrWithCode = err !== null && typeof err === 'object' && 'code' in err;
        const code = isErrWithCode ? (err as Record<string, unknown>).code : undefined;
        const errno = isErrWithCode ? (err as Record<string, unknown>).errno : undefined;
        // MySQL: ER_DUP_ENTRY (code 'ER_DUP_ENTRY') or errno 1062
        if (code === 'ER_DUP_ENTRY' || errno === 1062) {
          const existing = await this.repo.findOne({ where: { singletonKey: SINGLETON_KEY } });
          if (existing) return existing;
        }
        throw err;
      }
      return config;
    }

    // Migrate old defaults: if provider is still OPENAI (old default) and config was never
    // explicitly configured (no API key, never verified), switch to DeepSeek silently.
    if (
      config.provider === AiProvider.OPENAI &&
      config.baseUrl === DEFAULT_BASE_URLS[AiProvider.OPENAI] &&
      !config.encryptedApiKey &&
      !config.verified
    ) {
      config.provider = AiProvider.DEEPSEEK;
      config.baseUrl = DEFAULT_BASE_URLS[AiProvider.DEEPSEEK];
      await this.repo.save(config);
    }

    return config;
  }

  /** Build masked key display */
  private buildMaskedKey(keyLast4: string | null): string | null {
    if (keyLast4 && keyLast4.length === 4) {
      return `••••${keyLast4}`;
    }
    return null;
  }

  /** GET response */
  async getConfig(): Promise<AiConfigResponseDto> {
    const config = await this.getOrCreateConfig();
    const { source } = this.resolveApiKey(config);

    const hasDbKey = !!(config.encryptedApiKey && config.apiKeyIv && config.apiKeyAuthTag);

    return {
      id: config.id,
      provider: config.provider,
      baseUrl: config.baseUrl,
      hasApiKey: source !== 'none',
      hasDatabaseKey: hasDbKey,
      maskedApiKey: config.keyLast4
        ? this.buildMaskedKey(config.keyLast4)
        : source !== 'none'
          ? '••••'
          : null,
      keySource: source,
      defaultModel: config.defaultModel ?? null,
      enabled: config.enabled,
      supportsVision: config.supportsVision,
      timeoutMs: config.timeoutMs,
      reasoningEffort: config.reasoningEffort ?? null,
      verified: config.verified,
      lastTestedAt: config.lastTestedAt?.toISOString() ?? null,
      lastTestLatencyMs: config.lastTestLatencyMs ?? null,
      createdAt: config.createdAt.toISOString(),
      updatedAt: config.updatedAt.toISOString(),
    };
  }

  /** PUT / save */
  async saveConfig(dto: SaveAiConfigDto): Promise<AiConfig> {
    const config = await this.getOrCreateConfig();

    const normalizedBaseUrl = validateAndNormalizeBaseUrl(dto.baseUrl, dto.provider);

    // DNS SSRF check for all providers
    await validateDnsNotPrivate(new URL(normalizedBaseUrl).hostname);

    config.provider = dto.provider;
    config.baseUrl = normalizedBaseUrl;

    if (dto.defaultModel !== undefined) {
      config.defaultModel = dto.defaultModel || null;
    }

    if (dto.timeoutMs !== undefined) {
      config.timeoutMs = dto.timeoutMs;
    }

    if (dto.reasoningEffort !== undefined) {
      config.reasoningEffort = dto.reasoningEffort || null;
    }

    // Handle apiKey — empty/undefined = keep existing
    if (dto.apiKey !== undefined && dto.apiKey !== '') {
      const { ciphertext, iv, authTag } = encrypt(dto.apiKey);
      config.encryptedApiKey = ciphertext;
      config.apiKeyIv = iv;
      config.apiKeyAuthTag = authTag;
      config.keyLast4 = dto.apiKey.slice(-4);
    }

    // AI is always enabled by default — the enable switch has been removed
    if (dto.enabled !== undefined) {
      config.enabled = dto.enabled;
    } else {
      config.enabled = true;
    }

    if (dto.supportsVision !== undefined) {
      config.supportsVision = dto.supportsVision;
    }

    if (dto.enabled === true) {
      const { plaintext } = this.resolveApiKey(config);
      if (!plaintext) throw new BadRequestException('启用 AI 服务前必须配置 API Key');
      if (!config.defaultModel?.trim()) {
        throw new BadRequestException('启用 AI 服务前必须配置默认模型');
      }
    }

    return this.repo.save(config);
  }

  /** Clear DB key only */
  async clearKey(): Promise<AiConfigResponseDto> {
    const config = await this.getOrCreateConfig();
    config.encryptedApiKey = null;
    config.apiKeyIv = null;
    config.apiKeyAuthTag = null;
    config.keyLast4 = null;
    // If no env key either, disable
    const envKey = process.env.AI_API_KEY;
    if (!envKey) {
      config.enabled = false;
    }
    await this.repo.save(config);

    return this.getConfig();
  }

  /** Test connection — uses saved config or request body overrides */
  async testConnection(dto?: TestAiConfigDto): Promise<AiConfigTestResultDto> {
    return testConnection(this, dto);
  }

  /** Fetch available model list from the configured provider */
  async fetchModels(dto?: FetchModelsDto): Promise<FetchModelsResultDto> {
    return fetchModels(this, dto);
  }

  /**
   * Server-only runtime config — for future AI adapters.
   * Re-validates the stored base URL and DNS at runtime to guard
   * against config-table tampering or DNS record changes.
   * Future adapters should still use a restricted transport helper.
   */
  async getRuntimeConfig(): Promise<AiRuntimeConfig> {
    const config = await this.repo.findOne({ where: { singletonKey: SINGLETON_KEY } });
    if (!config) {
      throw new InternalServerErrorException('AI 配置未初始化');
    }

    if (!config.enabled) {
      throw new BadRequestException('AI 服务未启用');
    }

    // Re-validate and normalize the stored base URL
    const normalizedBaseUrl = validateAndNormalizeBaseUrl(config.baseUrl, config.provider);

    // Re-check DNS at runtime
    await validateDnsNotPrivate(new URL(normalizedBaseUrl).hostname);

    const { plaintext } = this.resolveApiKey(config);
    if (!plaintext) {
      throw new BadRequestException('未配置 API Key');
    }

    // defaultModel is required for actual AI calls
    if (!config.defaultModel) {
      throw new BadRequestException('未配置默认模型');
    }

    return {
      provider: config.provider,
      baseUrl: normalizedBaseUrl,
      apiKey: plaintext,
      defaultModel: config.defaultModel,
      timeoutMs: config.timeoutMs,
      enabled: config.enabled,
      supportsVision: config.supportsVision,
      reasoningEffort: config.reasoningEffort ?? null,
    };
  }
}