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 { estimateTeacherTokens, type TeacherModelMessage, type TeacherToolCall } from './context';
import { createTeacherReadTools, type TeacherReadAccess, type TeacherReadTools } from './read-tools';

interface TeacherModelConfig {
  api_key: string;
  base_url: string;
  models: string[];
  model_capabilities_v2: unknown;
}
export async function prepareTeacherModel(account: TeacherAccount, definition: TeacherDefinition, access?: TeacherReadAccess) {
  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
  );
  const tools = access ? createTeacherReadTools(access) : undefined;
  const toolBudget = tools ? Buffer.byteLength(JSON.stringify(tools.definitions), 'utf8') + 64 : 0;
  return {
    // Leave room for a read result and its native tool-call envelope.
    inputLimit: inputLimit - toolBudget - (tools ? Math.min(2400, Math.floor(inputLimit / 4)) : 0),
    run: (messages: TeacherModelMessage[], signal: AbortSignal, onText: (text: string) => void) => {
      assertTeacherAccount(account);
      return streamTeacherReply(config, modelId, fields, outputLimit, messages, signal, onText,
        proxyAwareFetch, { tools, inputLimit, assertCurrent: () => assertTeacherAccount(account) });
    },
  };
}
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,
  options?: { tools?: TeacherReadTools; inputLimit: number; assertCurrent(): void }
): Promise<PublicUsage | undefined> {
  const tools = options?.tools;
  const toolBudget = tools ? Buffer.byteLength(JSON.stringify(tools.definitions), 'utf8') + 64 : 0;
  const reads: TeacherModelMessage[][] = [];
  let usage: PublicUsage | undefined;
  // Six read rounds, then one final text response. No recursive agent or Pi session.
  for (let round = 0; round <= 6; round++) {
    signal.throwIfAborted();
    options?.assertCurrent();
    const build = () => [...messages, ...reads.flat()];
    while (options && estimateTeacherTokens(build()) + toolBudget > options.inputLimit && reads.length > 1) {
      reads.shift();
    }
    if (options && estimateTeacherTokens(build()) + toolBudget > options.inputLimit) {
      throw new TeacherError(422, 'teacher_context_too_long', '老师读取的内容超过上下文预算，请缩小问题范围或联系运营增加预算。');
    }
    const result = await streamTeacherTurn(config, modelId, reasoningFields, outputLimit,
      build(), signal, onText, fetchImpl, tools, round === 6);
    if (result.usage) usage = {
      inputTokens: (usage?.inputTokens ?? 0) + result.usage.inputTokens,
      outputTokens: (usage?.outputTokens ?? 0) + result.usage.outputTokens,
    };
    if (!result.calls.length) return usage;
    if (!tools || round === 6) {
      throw new TeacherError(502, 'teacher_tools_unavailable', '老师未能完成本次读取，请缩小问题范围后重试。');
    }
    const batch: TeacherModelMessage[] = [{ role: 'assistant', content: result.text, tool_calls: result.calls,
      ...(result.reasoning ? { reasoning_content: result.reasoning } : {}) }];
    const resultBudget = Math.min(2400, Math.floor(((options?.inputLimit ?? Infinity)
      - toolBudget - estimateTeacherTokens([...messages, ...batch]) - 64 * result.calls.length) / result.calls.length));
    if (resultBudget < 128) {
      throw new TeacherError(422, 'teacher_context_too_long', '老师读取的内容超过上下文预算，请缩小问题范围或联系运营增加预算。');
    }
    for (const call of result.calls) {
      signal.throwIfAborted();
      options?.assertCurrent();
      batch.push({ role: 'tool', tool_call_id: call.id,
        content: await tools.execute(call.function.name, call.function.arguments, signal, resultBudget) });
    }
    reads.push(batch);
  }
  return usage;
}

async function streamTeacherTurn(
  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>,
  tools: TeacherReadTools | undefined,
  finalRound: boolean
) {
  // 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,
      ...(tools ? { tools: tools.definitions, tool_choice: finalRound ? 'none' : 'auto', parallel_tool_calls: false } : {}),
    }),
    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;
  let text = '', reasoning = '', finishReason: string | null = null;
  const calls = new Map<number, TeacherToolCall>();
  const frame = (data: string) => {
    if (data === '[DONE]') {
      settled = true;
      return;
    }
    const event = JSON.parse(data) as {
      error?: unknown;
      choices?: Array<{
        delta?: { content?: unknown; reasoning_content?: unknown; tool_calls?: Array<{
          index: number; id?: string; function?: { name?: string; arguments?: string };
        }> };
        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 (!tools && (choice?.delta?.tool_calls || choice?.finish_reason === 'tool_calls'))
      throw new TeacherError(
        502,
        'teacher_tools_unavailable',
        '老师只能提供文字建议，本次回复未完成。'
      );
    if (typeof choice?.delta?.content === 'string') {
      text += choice.delta.content;
      onText(choice.delta.content);
    }
    // Some native reasoning providers require this on the next tool round. It stays Main-private.
    if (typeof choice?.delta?.reasoning_content === 'string') reasoning += choice.delta.reasoning_content;
    for (const delta of choice?.delta?.tool_calls ?? []) {
      if (!Number.isSafeInteger(delta.index) || delta.index < 0 || delta.index >= 8)
        throw new TeacherError(502, 'teacher_stream_invalid', '老师读取请求格式无效。');
      const call = calls.get(delta.index) ?? { id: '', type: 'function', function: { name: '', arguments: '' } };
      if (delta.id) call.id = delta.id;
      if (delta.function?.name) call.function.name += delta.function.name;
      if (delta.function?.arguments) call.function.arguments += delta.function.arguments;
      calls.set(delta.index, call);
    }
    if (choice?.finish_reason) finishReason = choice.finish_reason;
    if (['stop', 'length', 'tool_calls'].includes(finishReason ?? '')) 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', '回复中断，以下内容可能不完整。');
    if ((calls.size && finishReason !== 'tool_calls') || (finishReason === 'tool_calls' && !calls.size)
      || [...calls.values()].some(call => !call.id || !call.function.name))
      throw new TeacherError(502, 'teacher_stream_interrupted', '老师读取请求未完整收到，请重试。');
    return { usage, text, reasoning, calls: [...calls.values()] };
  } finally {
    await reader.cancel().catch(() => undefined);
    reader.releaseLock();
  }
}