import {
  AiStructuredOutputClient,
  type TokenUsage,
} from '../ai-shared/AiStructuredOutputClient';
import { getAiSemaphore } from '../ai-shared/AiConcurrency';
import { AiApiError, AiErrorCodes } from '../ai-shared/AiError';
import { AI_LIMITS, utf8ByteLength } from '../ai-shared/AiLimits';
import { resolveAiRuntimeConfig, type AiRuntimeConfig } from '../ai-shared/AiRuntimeConfig';
import type { AiSemaphore } from '../ai-shared/AiSemaphore';
import { createAiTimeoutBudget } from '../ai-shared/AiTimeout';
import { AiApiRegistry } from './AiApiRegistry';
import { AI_FUNC_PATTERN } from './AiApiDslParser';
import { buildAiFunctionPrompt } from './AiPromptBuilder';

export interface AiInvokeOptions {
  /**
   * 调用方已持有上游并发槽位时跳过再次申请。
   *
   * 助手流在整个 streamText 期间占着全局信号量的一个槽位,而本方法内部申请的是
   * 同一个模块级单例信号量;桥接工具若直接嵌套调用,满并发时每一路助手都会等一个
   * 永远不会释放的槽位,全体阻塞到超时。槽位继承把不变量钉死为"1 路助手 = 1 上游槽位"。
   */
  withinAiSlot?: boolean;
}

export interface AiInvokeResult {
  output: Record<string, string>;
  meta: {
    model: string;
    durationMs: number;
    retried: boolean;
    usage: TokenUsage;
  };
}

interface RuntimeState {
  config: AiRuntimeConfig;
  registry: AiApiRegistry;
  client: AiStructuredOutputClient;
  semaphore: AiSemaphore;
}

export class AiApiRuntime {
  private static instance: AiApiRuntime | undefined;
  private statePromise: Promise<RuntimeState> | undefined;

  static getInstance(): AiApiRuntime {
    return this.instance ??= new AiApiRuntime();
  }

  async invoke(
    func: string,
    input: Record<string, string>,
    signal?: AbortSignal,
    options?: AiInvokeOptions,
  ): Promise<AiInvokeResult> {
    const startedAt = performance.now();
    if (typeof func !== 'string' || !AI_FUNC_PATTERN.test(func)) {
      throw new AiApiError(AiErrorCodes.AI_INPUT_INVALID, 'func 格式不合法');
    }
    const state = await this.getState();
    const definition = state.registry.get(func);
    if (!definition) {
      throw new AiApiError(AiErrorCodes.AI_FUNC_NOT_FOUND, `AI 函数未注册：${func}`);
    }
    validateInput(func, input, definition.input);
    const timeoutMs = definition.options.timeoutMs ?? state.config.timeoutMs;
    // 嵌套调用的预算天然取 min(func timeoutMs, 调用方剩余预算):budget 由父 signal 派生。
    const budget = createAiTimeoutBudget(timeoutMs, signal);
    const runWithSlot = options?.withinAiSlot
      ? <T>(task: () => Promise<T>) => task()
      : <T>(task: () => Promise<T>) => state.semaphore.run(task, budget.signal);
    try {
      return await runWithSlot(async () => {
        budget.throwIfExpired();
        const prompt = buildAiFunctionPrompt(definition, input);
        const model = definition.options.model ?? state.config.aiApi.model ?? state.config.model;
        const maxOutputTokens = definition.options.maxOutputTokens
          ?? state.config.aiApi.maxOutputTokens;
        let retried = false;
        let usage: TokenUsage = { inputTokens: 0, outputTokens: 0 };
        for (let attempt = 0; ; attempt += 1) {
          budget.throwIfExpired();
          try {
            const response = await state.client.generateObject({
              ...prompt,
              schema: definition.outputSchema,
              model,
              maxOutputTokens,
              timeoutMs: budget.remainingMs(),
              signal: budget.signal,
            });
            usage = addUsage(usage, response.usage);
            const output = validateOutput(response.object, definition.output);
            if (output) {
              return {
                output,
                meta: {
                  model: response.model,
                  durationMs: Math.round(performance.now() - startedAt),
                  retried,
                  usage,
                },
              };
            }
          } catch (error) {
            if (!(error instanceof AiApiError) || error.code !== AiErrorCodes.AI_OUTPUT_INVALID) {
              throw error;
            }
          }
          if (attempt >= state.config.aiApi.retryOnParseError) {
            throw new AiApiError(AiErrorCodes.AI_OUTPUT_INVALID, 'AI 输出不符合函数契约');
          }
          retried = true;
        }
      });
    } catch (error) {
      if (budget.signal.aborted && !(error instanceof AiApiError)) {
        throw new AiApiError(AiErrorCodes.AI_UPSTREAM_TIMEOUT, 'AI 服务调用超时');
      }
      throw error;
    } finally {
      budget.dispose();
    }
  }

  async list(): Promise<ReturnType<AiApiRegistry['list']>> {
    return (await this.getState()).registry.list();
  }

  async listAssistantTools(): Promise<ReturnType<AiApiRegistry['listAssistantTools']>> {
    return (await this.getState()).registry.listAssistantTools();
  }

  private getState(): Promise<RuntimeState> {
    return this.statePromise ??= Promise.resolve().then(() => {
      const config = resolveAiRuntimeConfig();
      return {
        config,
        registry: AiApiRegistry.compile(config),
        client: new AiStructuredOutputClient(config),
        semaphore: getAiSemaphore(config.maxConcurrency),
      };
    }).catch((error) => {
      this.statePromise = undefined;
      throw error;
    });
  }
}

export function validateInput(
  func: unknown,
  input: unknown,
  expected: Readonly<Record<string, string>>,
): asserts input is Record<string, string> {
  if (typeof func !== 'string' || !AI_FUNC_PATTERN.test(func)) {
    throw new AiApiError(AiErrorCodes.AI_INPUT_INVALID, 'func 格式不合法');
  }
  if (!input || typeof input !== 'object' || Array.isArray(input)) {
    throw new AiApiError(AiErrorCodes.AI_INPUT_INVALID, 'input 必须是字符串对象');
  }
  const values = input as Record<string, unknown>;
  const expectedKeys = Object.keys(expected);
  const actualKeys = Object.keys(values);
  const missing = expectedKeys.filter((key) => !(key in values));
  const extra = actualKeys.filter((key) => !(key in expected));
  const invalid = actualKeys.filter((key) => typeof values[key] !== 'string');
  if (missing.length || extra.length || invalid.length) {
    const details = [
      missing.length ? `缺少变量：${missing.join(', ')}` : '',
      extra.length ? `多余变量：${extra.join(', ')}` : '',
      invalid.length ? `非字符串变量：${invalid.join(', ')}` : '',
    ].filter(Boolean).join('；');
    throw new AiApiError(AiErrorCodes.AI_INPUT_INVALID, details);
  }
  let totalBytes = 0;
  for (const [name, value] of Object.entries(values)) {
    const bytes = utf8ByteLength(value as string);
    if (bytes > AI_LIMITS.invokeInputValueBytes) {
      throw new AiApiError(AiErrorCodes.AI_INPUT_INVALID, `输入变量 ${name} 超过 64KB`);
    }
    totalBytes += bytes;
  }
  if (totalBytes > AI_LIMITS.invokeInputTotalBytes) {
    throw new AiApiError(AiErrorCodes.AI_INPUT_INVALID, '输入总大小超过 256KB');
  }
}

function validateOutput(
  value: Record<string, unknown>,
  expected: Readonly<Record<string, string>>,
): Record<string, string> | null {
  const expectedKeys = Object.keys(expected).sort();
  const actualKeys = Object.keys(value).sort();
  if (
    expectedKeys.length !== actualKeys.length
    || expectedKeys.some((key, index) => key !== actualKeys[index])
    || actualKeys.some((key) => typeof value[key] !== 'string')
  ) {
    return null;
  }
  return value as Record<string, string>;
}

function addUsage(left: TokenUsage, right: TokenUsage): TokenUsage {
  return {
    inputTokens: left.inputTokens + right.inputTokens,
    outputTokens: left.outputTokens + right.outputTokens,
  };
}