import { proxyAwareFetch } from '../utils/proxy-fetch';
import {
  buildManagedModelRequest,
  normalizeManagedModelCatalog,
} from '../../shared/managed-model-capabilities';
import type { PublicUsage } from '../../shared/coding-conversation-contracts';
import type { TeacherDefinition } from '../../shared/coding-teacher';
import {
  assertTeacherAccount,
  teacherCloudRequest,
  TeacherError,
  type TeacherAccount,
} from './config-client';
import type { TeacherModelMessage } from './context';

interface TeacherModelConfig {
  api_key: string;
  base_url: string;
  models: string[];
  model_capabilities_v2: unknown;
}
export async function prepareTeacherModel(account: TeacherAccount, definition: TeacherDefinition) {
  const config = await teacherCloudRequest<TeacherModelConfig>(
    account,
    '/api/auth/me/model-config'
  );
  const modelId = definition.model.model_id;
  const catalog = normalizeManagedModelCatalog(config.model_capabilities_v2, config.models);
  const capability = modelId ? catalog?.models[modelId] : null;
  if (
    !modelId ||
    !config.models.includes(modelId) ||
    !capability?.inputModalities?.includes('text') ||
    !capability.outputModalities?.includes('text')
  ) {
    throw new TeacherError(
      422,
      'teacher_model_unavailable',
      '老师所用模型暂不可用，请联系运营调整。'
    );
  }
  const savedChoice = definition.model.reasoning_choice;
  // Published definitions serialize an unspecified effort as null.
  const choice = savedChoice.mode === 'enabled'
    ? { mode: savedChoice.mode, ...(savedChoice.effort == null ? {} : { effort: savedChoice.effort }) }
    : { mode: savedChoice.mode };
  let fields: Record<string, unknown>;
  try {
    fields = buildManagedModelRequest(
      modelId,
      choice,
      capability
    ).reasoningFields;
  } catch {
    throw new TeacherError(
      422,
      'teacher_model_unavailable',
      '老师所用思考选项已不可用，请联系运营调整。'
    );
  }
  const outputLimit = Math.min(
    definition.limits.max_output_tokens,
    capability.limits?.maxOutputTokens ?? Infinity
  );
  const inputLimit = Math.min(
    definition.limits.max_input_tokens,
    capability.limits?.maxInputTokens ?? Infinity,
    capability.limits?.contextWindow ? capability.limits.contextWindow - outputLimit : Infinity
  );
  return {
    inputLimit,
    run: (messages: TeacherModelMessage[], signal: AbortSignal, onText: (text: string) => void) => {
      assertTeacherAccount(account);
      return streamTeacherReply(config, modelId, fields, outputLimit, messages, signal, onText);
    },
  };
}
export async function streamTeacherReply(
  config: Pick<TeacherModelConfig, 'api_key' | 'base_url'>,
  modelId: string,
  reasoningFields: Record<string, unknown>,
  outputLimit: number,
  messages: TeacherModelMessage[],
  signal: AbortSignal,
  onText: (text: string) => void,
  fetchImpl: (input: string | URL, init?: RequestInit) => Promise<Response> = proxyAwareFetch
): Promise<PublicUsage | undefined> {
  // The gateway base already includes its version prefix, as in the existing AI proxy.
  const response = await fetchImpl(config.base_url.replace(/\/+$/, '') + '/chat/completions', {
    method: 'POST',
    headers: { 'Content-Type': 'application/json', 'x-works-square-ai-token': config.api_key },
    body: JSON.stringify({
      model: modelId,
      messages,
      stream: true,
      stream_options: { include_usage: true },
      max_tokens: outputLimit,
      ...reasoningFields,
    }),
    signal,
  });
  if (!response.ok) {
    throw new TeacherError(
      response.status,
      'teacher_model_failed',
      response.status === 402
        ? '词元点数不足，请补充后再提问。'
        : response.status === 401
          ? '模型凭据已失效，请重新提问。'
          : '老师暂时无法回复，请稍后重试。'
    );
  }
  if (!response.body) throw new TeacherError(502, 'teacher_stream_missing', '未收到老师回复。');
  const reader = response.body.getReader(),
    decoder = new TextDecoder();
  let buffer = '',
    settled = false,
    usage: PublicUsage | undefined;
  const frame = (data: string) => {
    if (data === '[DONE]') {
      settled = true;
      return;
    }
    const event = JSON.parse(data) as {
      error?: unknown;
      choices?: Array<{
        delta?: { content?: unknown; tool_calls?: unknown };
        finish_reason?: string | null;
      }>;
      usage?: { prompt_tokens: number; completion_tokens: number };
    };
    if (event.error)
      throw new TeacherError(502, 'teacher_model_failed', '老师回复中断，请保留当前内容后重试。');
    const choice = event.choices?.[0];
    if (choice?.delta?.tool_calls || choice?.finish_reason === 'tool_calls')
      throw new TeacherError(
        502,
        'teacher_tools_unavailable',
        '老师只能提供文字建议，本次回复未完成。'
      );
    if (typeof choice?.delta?.content === 'string') onText(choice.delta.content);
    if (choice?.finish_reason === 'stop' || choice?.finish_reason === 'length') settled = true;
    if (
      event.usage &&
      Number.isFinite(event.usage.prompt_tokens) &&
      Number.isFinite(event.usage.completion_tokens)
    )
      usage = {
        inputTokens: event.usage.prompt_tokens,
        outputTokens: event.usage.completion_tokens,
      };
  };
  try {
    while (true) {
      const chunk = await reader.read();
      buffer += decoder.decode(chunk.value, { stream: !chunk.done }).replace(/\r\n/g, '\n');
      let boundary: number;
      while ((boundary = buffer.indexOf('\n\n')) >= 0) {
        const event = buffer.slice(0, boundary);
        buffer = buffer.slice(boundary + 2);
        const data = event
          .split('\n')
          .filter((line) => line.startsWith('data:'))
          .map((line) => line.slice(5).trimStart())
          .join('\n');
        if (data) frame(data);
      }
      if (buffer.length > 1_048_576)
        throw new TeacherError(502, 'teacher_stream_invalid', '老师回复格式无效。');
      if (chunk.done) break;
    }
    if (signal.aborted) throw signal.reason;
    if (!settled)
      throw new TeacherError(502, 'teacher_stream_interrupted', '回复中断，以下内容可能不完整。');
    return usage;
  } finally {
    await reader.cancel().catch(() => undefined);
    reader.releaseLock();
  }
}