import { BadGatewayException, Injectable, RequestTimeoutException } from '@nestjs/common';
import { lookup } from 'node:dns';
import * as http from 'node:http';
import * as https from 'node:https';
import { isIP } from 'node:net';
import type { AiRuntimeConfig } from '../ai-config/dto/ai-config.dto';
import { AiProvider } from '../ai-config/ai-config.entity';
import type { ModelMessage, ModelStreamEvent } from './ai-chat.types';

interface ChatTool {
  type: 'function';
  function: {
    name: string;
    description: string;
    parameters: Record<string, unknown>;
  };
}

interface StreamChoiceDelta {
  content?: string | null;
  reasoning_content?: string | null;
  tool_calls?: Array<{
    index?: number;
    id?: string;
    function?: { name?: string; arguments?: string };
  }>;
}

const MAX_UPSTREAM_EVENT_BYTES = 1024 * 1024;
const MAX_UPSTREAM_RETRIES = 3;
const UPSTREAM_RETRY_DELAYS_MS = [500, 1000, 2000];
const RETRYABLE_STATUS_CODES = new Set([408, 425, 429, 500, 502, 503, 504]);
const RETRYABLE_TRANSPORT_CODES = new Set([
  'ECONNRESET',
  'ECONNREFUSED',
  'ETIMEDOUT',
  'ENOTFOUND',
  'EAI_AGAIN',
  'EPIPE',
]);

// Known public provider hosts — trusted even if CDN resolves to private-range IPs
const DNS_TRUSTED_HOSTS = new Set(['api.openai.com', 'api.deepseek.com']);

const PRIVATE_IPV4_RANGES = [
  /^127\./,
  /^10\./,
  /^172\.(1[6-9]|2\d|3[01])\./,
  /^192\.168\./,
  /^169\.254\./,
  /^0\./,
  /^100\.(6[4-9]|[7-9]\d|1[01]\d|12[0-7])\./,
];

interface PinnedResponse {
  status: number;
  contentType: string;
  body: http.IncomingMessage;
}

@Injectable()
export class AiModelStreamService {
  async *stream(
    config: AiRuntimeConfig,
    messages: ModelMessage[],
    tools: ChatTool[],
    signal: AbortSignal,
  ): AsyncGenerator<ModelStreamEvent> {
    const requestBody = JSON.stringify({
      model: config.defaultModel,
      messages,
      stream: true,
      ...(tools.length ? { tools, tool_choice: 'auto' } : {}),
      // reasoning_effort 仅对支持该参数的 OpenAI 兼容服务生效；
      // DeepSeek 官方接口不接受该参数，避免请求被拒。
      ...(config.reasoningEffort &&
      config.reasoningEffort !== 'none' &&
      config.provider !== AiProvider.DEEPSEEK
        ? { reasoning_effort: config.reasoningEffort }
        : {}),
    });
    const url = `${config.baseUrl.replace(/\/$/, '')}/chat/completions`;
    const headers = {
      Authorization: `Bearer ${config.apiKey}`,
      'Content-Type': 'application/json',
      Accept: 'text/event-stream',
    };
    let response: PinnedResponse | null = null;
    let activeTimeout: AbortSignal | undefined;

    for (let attempt = 0; attempt <= MAX_UPSTREAM_RETRIES; attempt += 1) {
      activeTimeout = AbortSignal.timeout(config.timeoutMs);
      const combinedSignal = AbortSignal.any([signal, activeTimeout]);
      try {
        response = await this.pinnedPost(url, headers, requestBody, combinedSignal);
      } catch (error) {
        if (activeTimeout.aborted && !signal.aborted) {
          throw new RequestTimeoutException('AI 服务响应超时');
        }
        if (
          attempt < MAX_UPSTREAM_RETRIES &&
          !signal.aborted &&
          this.isRetryableTransportError(error)
        ) {
          const delayMs = UPSTREAM_RETRY_DELAYS_MS[attempt];
          yield {
            type: 'retrying',
            attempt: attempt + 1,
            maxRetries: MAX_UPSTREAM_RETRIES,
            delayMs,
            reason: error instanceof Error ? error.message : '网络连接失败',
          };
          await this.sleep(delayMs);
          continue;
        }
        throw error;
      }

      if (response.status >= 200 && response.status < 300) break;
      if (attempt < MAX_UPSTREAM_RETRIES && RETRYABLE_STATUS_CODES.has(response.status)) {
        response.body.resume?.();
        const delayMs = UPSTREAM_RETRY_DELAYS_MS[attempt];
        yield {
          type: 'retrying',
          attempt: attempt + 1,
          maxRetries: MAX_UPSTREAM_RETRIES,
          delayMs,
          reason: `上游返回 ${response.status}`,
        };
        await this.sleep(delayMs);
        continue;
      }
      const body = await this.readLimitedBody(response.body);
      throw new BadGatewayException(this.safeUpstreamMessage(response.status, body));
    }

    if (!response) throw new BadGatewayException('AI 服务暂时不可用');

    if (response.status < 200 || response.status >= 300) {
      const body = await this.readLimitedBody(response.body);
      throw new BadGatewayException(this.safeUpstreamMessage(response.status, body));
    }
    const contentType = response.contentType.toLowerCase();
    if (!contentType.includes('text/event-stream')) {
      throw new BadGatewayException('AI 服务返回了无效的响应格式');
    }
    const decoder = new TextDecoder();
    let buffer = '';
    const calls = new Map<number, { id: string; name: string; arguments: string }>();

    const consumeEvent = (event: string): ModelStreamEvent[] => {
      const output: ModelStreamEvent[] = [];
      const data = event
        .split(/\r?\n/)
        .filter((line) => line.startsWith('data:'))
        .map((line) => line.slice(5).trimStart())
        .join('\n');
      if (!data || data === '[DONE]') return output;
      const parsed = this.parseEvent(data);
      const delta = parsed.choices?.[0]?.delta;
      if (!delta) return output;
      if (delta.reasoning_content) output.push({ type: 'reasoning', delta: delta.reasoning_content });
      if (delta.content) output.push({ type: 'content', delta: delta.content });
      for (const part of delta.tool_calls ?? []) {
        const index = part.index ?? 0;
        const current = calls.get(index) ?? { id: '', name: '', arguments: '' };
        if (part.id) current.id += part.id;
        if (part.function?.name) current.name += part.function.name;
        if (part.function?.arguments) current.arguments += part.function.arguments;
        calls.set(index, current);
      }
      return output;
    };

    try {
      for await (const chunk of response.body as AsyncIterable<Uint8Array>) {
        buffer += decoder.decode(chunk, { stream: true });
        if (buffer.length > MAX_UPSTREAM_EVENT_BYTES) {
          throw new BadGatewayException('AI 服务返回的单个事件过大');
        }
        const events = buffer.split(/\r?\n\r?\n/);
        buffer = events.pop() ?? '';
        for (const event of events) for (const parsed of consumeEvent(event)) yield parsed;
      }
      buffer += decoder.decode();
      if (buffer.trim()) for (const parsed of consumeEvent(buffer)) yield parsed;
    } catch (error) {
      if (activeTimeout?.aborted && !signal.aborted) {
        throw new RequestTimeoutException('AI 服务响应超时');
      }
      throw error;
    }

    yield {
      type: 'complete',
      toolCalls: [...calls.entries()]
        .sort(([a], [b]) => a - b)
        .map(([, call], index) => ({
          id: call.id || `call_${index}`,
          name: call.name,
          arguments: call.arguments || '{}',
        })),
    };
  }

  private parseEvent(data: string): { choices?: Array<{ delta?: StreamChoiceDelta }> } {
    try {
      const value: unknown = JSON.parse(data);
      if (!value || typeof value !== 'object') throw new Error('invalid');
      return value;
    } catch {
      throw new BadGatewayException('AI 服务返回了无效的流式数据');
    }
  }

  private isRetryableTransportError(error: unknown): boolean {
    if (!(error instanceof Error)) return false;
    const code = (error as NodeJS.ErrnoException).code;
    if (code && RETRYABLE_TRANSPORT_CODES.has(code)) return true;
    return /socket hang up|ECONNRESET|ETIMEDOUT|EAI_AGAIN/i.test(error.message);
  }

  private sleep(ms: number): Promise<void> {
    return new Promise((resolve) => setTimeout(resolve, ms));
  }

  private safeUpstreamMessage(status: number, body: string): string {
    if (status === 401 || status === 403) return 'AI 服务认证失败';
    if (status === 429) return 'AI 服务请求过于频繁';
    if (status === 503) return 'AI 服务繁忙，请稍后重试';
    if (status >= 500) return 'AI 服务暂时不可用';
    const message = this.extractErrorMessage(body);
    return message ? `AI 服务请求失败：${message}` : `AI 服务请求失败（${status}）`;
  }

  private extractErrorMessage(body: string): string | null {
    try {
      const parsed = JSON.parse(body) as { error?: { message?: unknown } };
      const message = parsed.error?.message;
      return typeof message === 'string' ? message.slice(0, 200) : null;
    } catch {
      return null;
    }
  }

  private pinnedPost(
    url: string,
    headers: Record<string, string>,
    body: string,
    signal: AbortSignal,
  ): Promise<PinnedResponse> {
    return new Promise((resolve, reject) => {
      const parsed = new URL(url);
      const isHttps = parsed.protocol === 'https:';
      const port = parsed.port ? Number(parsed.port) : isHttps ? 443 : 80;
      lookup(parsed.hostname, { all: true, family: 0 }, (dnsError, addresses) => {
        if (dnsError || !addresses?.length) return reject(new Error('DNS 解析失败'));
        const allowPrivate =
          process.env.AI_ALLOW_PRIVATE_BASE_URL === 'true' ||
          DNS_TRUSTED_HOSTS.has(parsed.hostname);
        if (!allowPrivate && addresses.some(({ address }) => this.isPrivateAddress(address))) {
          return reject(new Error('域名解析到内网地址'));
        }
        const target = addresses[0];
        const transport = isHttps ? https : http;
        const request = transport.request(
          {
            hostname: target.address,
            port,
            path: parsed.pathname + parsed.search,
            method: 'POST',
            headers: {
              ...headers,
              Host: parsed.hostname,
              'Content-Length': Buffer.byteLength(body).toString(),
            },
            servername: isHttps ? parsed.hostname : undefined,
            rejectUnauthorized: isHttps,
            family: target.family === 6 ? 6 : 4,
            signal,
          },
          (response) => {
            const status = response.statusCode ?? 500;
            if (status >= 300 && status < 400) {
              response.resume();
              response.destroy();
              reject(new Error('禁止重定向'));
              return;
            }
            resolve({
              status,
              contentType: String(response.headers['content-type'] ?? ''),
              body: response,
            });
          },
        );
        request.once('error', reject);
        request.end(body);
      });
    });
  }

  private async readLimitedBody(body: http.IncomingMessage): Promise<string> {
    const chunks: Uint8Array[] = [];
    let total = 0;
    for await (const value of body as AsyncIterable<Uint8Array>) {
      total += value.length;
      if (total > MAX_UPSTREAM_EVENT_BYTES) {
        body.destroy();
        return '';
      }
      chunks.push(value);
    }
    return Buffer.concat(chunks).toString('utf8');
  }

  private isPrivateAddress(rawAddress: string): boolean {
    const address = rawAddress.toLowerCase();
    if (isIP(address) === 4) return PRIVATE_IPV4_RANGES.some((range) => range.test(address));
    if (isIP(address) !== 6) return true;
    if (address === '::1' || address === '::') return true;
    if (address.startsWith('fc') || address.startsWith('fd')) return true;
    if (/^fe[89ab]/.test(address)) return true;
    if (address.startsWith('::ffff:') && isIP(address.slice(7)) === 4) {
      return PRIVATE_IPV4_RANGES.some((range) => range.test(address.slice(7)));
    }
    return false;
  }
}
