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 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;

// 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 timeout = AbortSignal.timeout(config.timeoutMs);
    const combinedSignal = AbortSignal.any([signal, timeout]);
    let response: PinnedResponse;

    try {
      response = await this.pinnedPost(
        `${config.baseUrl.replace(/\/$/, '')}/chat/completions`,
        {
          Authorization: `Bearer ${config.apiKey}`,
          'Content-Type': 'application/json',
          Accept: 'text/event-stream',
        },
        JSON.stringify({
          model: config.defaultModel,
          messages,
          stream: true,
          ...(tools.length ? { tools, tool_choice: 'auto' } : {}),
        }),
        combinedSignal,
      );
    } catch (error) {
      if (timeout.aborted && !signal.aborted) throw new RequestTimeoutException('AI 服务响应超时');
      throw error;
    }

    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 (timeout.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 safeUpstreamMessage(status: number, body: string): string {
    if (status === 401 || status === 403) return 'AI 服务认证失败';
    if (status === 429) 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;
  }
}