import { EventEmitter } from 'node:events';
import type { IncomingMessage, ServerResponse } from 'http';
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
import { handleAiProxyRoutes } from '@electron/api/routes/ai-proxy';
import {
  clearWorksSquareAIGatewayCredential,
  seedWorksSquareAIGatewayCredential,
} from '@electron/services/works-square-ai-gateway';
import {
  clearWorksSquareSession,
  storeWorksSquareSession,
} from '@electron/services/works-square-session';

const loggerWarnMock = vi.hoisted(() => vi.fn());

vi.mock('@electron/utils/logger', () => ({
  logger: {
    debug: vi.fn(),
    error: vi.fn(),
    info: vi.fn(),
    warn: (...args: unknown[]) => loggerWarnMock(...args),
  },
}));

function createRequest(
  method: string,
  body?: unknown,
  headers: Record<string, string> = {},
): IncomingMessage {
  const req = new EventEmitter();
  const rawBody = body === undefined
    ? undefined
    : (typeof body === 'string' ? body : JSON.stringify(body));
  Object.assign(req, {
    method,
    headers: {
      host: '127.0.0.1:13210',
      authorization: 'Bearer local-host-api-token',
      'content-type': 'application/json',
      ...(rawBody ? { 'content-length': String(Buffer.byteLength(rawBody)) } : {}),
      ...headers,
    },
    [Symbol.asyncIterator]: async function* () {
      if (rawBody !== undefined) {
        yield Buffer.from(rawBody);
      }
    },
  });
  return req as IncomingMessage;
}

function createResponse(options: { writeResults?: boolean[] } = {}) {
  const chunks: string[] = [];
  const headers = new Map<string, number | string | string[]>();
  const writeResults = [...(options.writeResults ?? [])];
  const socket = { setNoDelay: vi.fn() };
  const res = new EventEmitter();
  Object.assign(res, {
    statusCode: 0,
    socket,
    flushHeaders: vi.fn(),
    setHeader: vi.fn((name: string, value: number | string | string[]) => {
      headers.set(name.toLowerCase(), value);
    }),
    write: vi.fn((chunk: string | Uint8Array) => {
      chunks.push(typeof chunk === 'string' ? chunk : Buffer.from(chunk).toString('utf8'));
      return writeResults.shift() ?? true;
    }),
    end: vi.fn((chunk?: string | Uint8Array) => {
      if (chunk !== undefined) {
        chunks.push(typeof chunk === 'string' ? chunk : Buffer.from(chunk).toString('utf8'));
      }
    }),
  });

  return {
    res: res as unknown as ServerResponse,
    socket,
    get statusCode() {
      return (res as { statusCode: number }).statusCode;
    },
    header: (name: string) => headers.get(name.toLowerCase()),
    body: () => chunks.join(''),
  };
}

function streamFromText(...chunks: string[]): ReadableStream<Uint8Array> {
  const encoder = new TextEncoder();
  return new ReadableStream<Uint8Array>({
    start(controller) {
      for (const chunk of chunks) {
        controller.enqueue(encoder.encode(chunk));
      }
      controller.close();
    },
  });
}

async function flushMicrotasks(): Promise<void> {
  for (let index = 0; index < 20; index += 1) {
    await Promise.resolve();
  }
}

describe('ai proxy routes', () => {
  beforeEach(() => {
    vi.useFakeTimers();
    vi.setSystemTime(new Date('2026-07-06T08:00:00.000Z'));
    loggerWarnMock.mockReset();
    clearWorksSquareAIGatewayCredential();
    clearWorksSquareSession();
  });

  afterEach(() => {
    vi.useRealTimers();
    clearWorksSquareAIGatewayCredential();
    clearWorksSquareSession();
    vi.unstubAllGlobals();
  });

  it('forwards OpenAI-compatible JSON requests to one-api with a fresh Works Square gateway token', async () => {
    seedWorksSquareAIGatewayCredential({
      accessToken: 'ws-ai-token',
      expiresIn: 3600,
      oneApiBaseUrl: 'https://one-api.example.com/v1',
    });
    const fetchMock = vi.fn().mockResolvedValueOnce(
      new Response(JSON.stringify({ id: 'chatcmpl_1' }), {
        status: 200,
        headers: { 'content-type': 'application/json' },
      }),
    );
    vi.stubGlobal('fetch', fetchMock);
    const response = createResponse();

    const handled = await handleAiProxyRoutes(
      createRequest('POST', { model: 'deepseek-chat', messages: [] }),
      response.res,
      new URL('http://127.0.0.1:13210/api/ai-proxy/v1/chat/completions?trace=1'),
      {} as never,
    );

    expect(handled).toBe(true);
    expect(fetchMock).toHaveBeenCalledWith(
      'https://one-api.example.com/v1/chat/completions?trace=1',
      expect.objectContaining({
        method: 'POST',
        body: Buffer.from(JSON.stringify({ model: 'deepseek-chat', messages: [] })),
        headers: expect.objectContaining({
          'content-type': 'application/json',
          'x-works-square-ai-token': 'ws-ai-token',
        }),
      }),
    );
    const forwardedHeaders = fetchMock.mock.calls[0][1].headers as Record<string, string>;
    expect(forwardedHeaders.authorization).toBeUndefined();
    expect(forwardedHeaders.host).toBeUndefined();
    expect(response.statusCode).toBe(200);
    expect(response.header('content-type')).toBe('application/json');
    expect(response.body()).toBe(JSON.stringify({ id: 'chatcmpl_1' }));
  });

  it('refreshes a near-expiry gateway credential before forwarding a long-running task request', async () => {
    storeWorksSquareSession({
      accessToken: 'works-access-token',
      refreshToken: 'works-refresh-token',
      expiresAt: Date.now() + 120_000,
    });
    seedWorksSquareAIGatewayCredential({
      accessToken: 'near-expiry-ai-token',
      expiresIn: 30,
      oneApiBaseUrl: 'https://one-api.example.com/v1',
    });
    const fetchMock = vi.fn()
      .mockResolvedValueOnce(new Response(JSON.stringify({
        access_token: 'fresh-ai-token',
        expires_in: 1800,
        one_api_base_url: 'https://one-api.example.com/v1',
      }), { status: 200 }))
      .mockResolvedValueOnce(new Response(JSON.stringify({ id: 'chatcmpl_refresh' }), {
        status: 200,
        headers: { 'content-type': 'application/json' },
      }));
    vi.stubGlobal('fetch', fetchMock);
    const response = createResponse();

    await handleAiProxyRoutes(
      createRequest('POST', { model: 'qwen3.7-max', messages: [] }),
      response.res,
      new URL('http://127.0.0.1:13210/api/ai-proxy/v1/chat/completions'),
      {} as never,
    );

    expect(fetchMock).toHaveBeenCalledTimes(2);
    expect(fetchMock.mock.calls[0][0]).toBe('https://square.nianxx.cn/api/ai-gateway/session');
    expect(fetchMock.mock.calls[1][1].headers).toMatchObject({
      'x-works-square-ai-token': 'fresh-ai-token',
    });
    expect(response.statusCode).toBe(200);
    expect(response.body()).toBe(JSON.stringify({ id: 'chatcmpl_refresh' }));
  });

  it('refreshes the gateway credential and retries once when one-api rejects an expired gateway token', async () => {
    storeWorksSquareSession({
      accessToken: 'works-access-token',
      refreshToken: 'works-refresh-token',
      expiresAt: Date.now() + 120_000,
    });
    seedWorksSquareAIGatewayCredential({
      accessToken: 'expired-ai-token',
      expiresIn: 3600,
      oneApiBaseUrl: 'https://one-api.example.com/v1',
    });
    const fetchMock = vi.fn()
      .mockResolvedValueOnce(new Response('AI access token expired', { status: 401 }))
      .mockResolvedValueOnce(
        new Response(JSON.stringify({
          access_token: 'fresh-ai-token',
          expires_in: 1800,
          one_api_base_url: 'https://one-api.example.com/v1',
        }), { status: 200 }),
      )
      .mockResolvedValueOnce(new Response(JSON.stringify({ id: 'chatcmpl_2' }), { status: 200 }));
    vi.stubGlobal('fetch', fetchMock);
    const response = createResponse();

    await handleAiProxyRoutes(
      createRequest('POST', { model: 'deepseek-chat' }),
      response.res,
      new URL('http://127.0.0.1:13210/api/ai-proxy/v1/chat/completions'),
      {} as never,
    );

    expect(fetchMock).toHaveBeenCalledTimes(3);
    expect(fetchMock.mock.calls[0][1].headers).toMatchObject({
      'x-works-square-ai-token': 'expired-ai-token',
    });
    expect(fetchMock.mock.calls[1][0]).toBe('https://square.nianxx.cn/api/ai-gateway/session');
    expect(fetchMock.mock.calls[2][1].headers).toMatchObject({
      'x-works-square-ai-token': 'fresh-ai-token',
    });
    expect(response.statusCode).toBe(200);
    expect(response.body()).toBe(JSON.stringify({ id: 'chatcmpl_2' }));
  });

  it('does not retry quota failures', async () => {
    seedWorksSquareAIGatewayCredential({
      accessToken: 'ws-ai-token',
      expiresIn: 3600,
      oneApiBaseUrl: 'https://one-api.example.com/v1',
    });
    const fetchMock = vi.fn().mockResolvedValueOnce(
      new Response(JSON.stringify({ error: 'user quota is not enough' }), {
        status: 429,
        headers: { 'content-type': 'application/json' },
      }),
    );
    vi.stubGlobal('fetch', fetchMock);
    const response = createResponse();

    await handleAiProxyRoutes(
      createRequest('POST', { model: 'deepseek-chat' }),
      response.res,
      new URL('http://127.0.0.1:13210/api/ai-proxy/v1/chat/completions'),
      {} as never,
    );

    expect(fetchMock).toHaveBeenCalledOnce();
    expect(response.statusCode).toBe(402);
    expect(response.body()).toBe(JSON.stringify({ error: 'user quota is not enough' }));
  });

  it('maps rolling-window quota exhaustion to a non-retryable response status', async () => {
    seedWorksSquareAIGatewayCredential({
      accessToken: 'ws-ai-token',
      expiresIn: 3600,
      oneApiBaseUrl: 'https://one-api.example.com/v1',
    });
    const fetchMock = vi.fn().mockResolvedValueOnce(
      new Response(JSON.stringify({
        error: {
          message: 'Token quota exhausted for rolling 5-hour window',
          code: 'works_square_gateway_authorize_failed',
          type: 'one_api_error',
        },
      }), {
        status: 429,
        headers: { 'content-type': 'application/json' },
      }),
    );
    vi.stubGlobal('fetch', fetchMock);
    const response = createResponse();

    await handleAiProxyRoutes(
      createRequest('POST', { model: 'qwen3.6-plus', messages: [] }),
      response.res,
      new URL('http://127.0.0.1:13210/api/ai-proxy/v1/chat/completions'),
      {} as never,
    );

    expect(fetchMock).toHaveBeenCalledOnce();
    expect(response.statusCode).toBe(402);
    expect(response.body()).toContain('Token quota exhausted for rolling 5-hour window');
  });

  it('preserves the upstream status for a generic gateway authorization failure', async () => {
    seedWorksSquareAIGatewayCredential({
      accessToken: 'ws-ai-token',
      expiresIn: 3600,
      oneApiBaseUrl: 'https://one-api.example.com/v1',
    });
    const fetchMock = vi.fn().mockResolvedValueOnce(
      new Response(JSON.stringify({
        error: {
          message: 'works_square_gateway_authorize_failed',
          code: 'works_square_gateway_authorize_failed',
          type: 'one_api_error',
        },
      }), {
        status: 429,
        headers: { 'content-type': 'application/json' },
      }),
    );
    vi.stubGlobal('fetch', fetchMock);
    const response = createResponse();

    await handleAiProxyRoutes(
      createRequest('POST', { model: 'qwen3.6-plus', messages: [] }),
      response.res,
      new URL('http://127.0.0.1:13210/api/ai-proxy/v1/chat/completions'),
      {} as never,
    );

    expect(fetchMock).toHaveBeenCalledOnce();
    expect(response.statusCode).toBe(429);
    expect(response.body()).toContain('works_square_gateway_authorize_failed');
  });

  it('logs a sanitized one-api error summary when the upstream group is saturated', async () => {
    seedWorksSquareAIGatewayCredential({
      accessToken: 'ws-ai-token',
      expiresIn: 3600,
      oneApiBaseUrl: 'https://one-api.example.com/v1',
    });
    const fetchMock = vi.fn().mockResolvedValueOnce(
      new Response(JSON.stringify({
        error: {
          message: '当前分组上游负载已饱和，请稍后再试 (request id: req-saturated)',
          code: 'rate_limit_exceeded',
          type: 'one_api_error',
        },
      }), {
        status: 429,
        headers: { 'content-type': 'application/json' },
      }),
    );
    vi.stubGlobal('fetch', fetchMock);
    const response = createResponse();

    await handleAiProxyRoutes(
      createRequest('POST', { model: 'qwen3.7-max', messages: [{ role: 'user', content: 'secret prompt' }] }),
      response.res,
      new URL('http://127.0.0.1:13210/api/ai-proxy/v1/chat/completions'),
      {} as never,
    );

    expect(response.statusCode).toBe(429);
    expect(response.body()).toContain('当前分组上游负载已饱和');
    expect(loggerWarnMock).toHaveBeenCalledWith(
      '[ai-proxy] One-api returned non-success response',
      expect.objectContaining({
        status: 429,
        url: 'https://one-api.example.com/v1/chat/completions',
        code: 'rate_limit_exceeded',
        type: 'one_api_error',
        message: '当前分组上游负载已饱和，请稍后再试 (request id: req-saturated)',
      }),
    );
    expect(JSON.stringify(loggerWarnMock.mock.calls)).not.toContain('secret prompt');
    expect(JSON.stringify(loggerWarnMock.mock.calls)).not.toContain('ws-ai-token');
  });

  it('strips decoded compression headers from proxied responses', async () => {
    seedWorksSquareAIGatewayCredential({
      accessToken: 'ws-ai-token',
      expiresIn: 3600,
      oneApiBaseUrl: 'https://one-api.example.com/v1',
    });
    const fetchMock = vi.fn().mockResolvedValueOnce(
      new Response(JSON.stringify({ id: 'chatcmpl_br' }), {
        status: 200,
        headers: {
          'content-encoding': 'br',
          'content-length': '999',
          'content-type': 'application/json',
        },
      }),
    );
    vi.stubGlobal('fetch', fetchMock);
    const response = createResponse();

    await handleAiProxyRoutes(
      createRequest('POST', { model: 'deepseek-chat' }, { 'accept-encoding': 'br, gzip' }),
      response.res,
      new URL('http://127.0.0.1:13210/api/ai-proxy/v1/chat/completions'),
      {} as never,
    );

    const forwardedHeaders = fetchMock.mock.calls[0][1].headers as Record<string, string>;
    expect(forwardedHeaders['accept-encoding']).toBeUndefined();
    expect(response.statusCode).toBe(200);
    expect(response.header('content-type')).toBe('application/json');
    expect(response.header('content-encoding')).toBeUndefined();
    expect(response.header('content-length')).toBeUndefined();
    expect(response.body()).toBe(JSON.stringify({ id: 'chatcmpl_br' }));
  });

  it('passes streaming chunks through to the local client', async () => {
    seedWorksSquareAIGatewayCredential({
      accessToken: 'ws-ai-token',
      expiresIn: 3600,
      oneApiBaseUrl: 'https://one-api.example.com/v1',
    });
    const fetchMock = vi.fn().mockResolvedValueOnce(
      new Response(streamFromText('data: one\n\n', 'data: two\n\n'), {
        status: 200,
        headers: { 'content-type': 'text/event-stream' },
      }),
    );
    vi.stubGlobal('fetch', fetchMock);
    const response = createResponse();

    await handleAiProxyRoutes(
      createRequest('POST', { stream: true }),
      response.res,
      new URL('http://127.0.0.1:13210/api/ai-proxy/v1/chat/completions'),
      {} as never,
    );

    expect(response.statusCode).toBe(200);
    expect(response.header('content-type')).toBe('text/event-stream');
    expect(response.body()).toBe('data: one\n\ndata: two\n\n');
  });

  it('flushes streaming headers and waits for downstream drain before reading more chunks', async () => {
    seedWorksSquareAIGatewayCredential({
      accessToken: 'ws-ai-token',
      expiresIn: 3600,
      oneApiBaseUrl: 'https://one-api.example.com/v1',
    });
    const fetchMock = vi.fn().mockResolvedValueOnce(
      new Response(streamFromText('data: one\n\n', 'data: two\n\n'), {
        status: 200,
        headers: { 'content-type': 'text/event-stream' },
      }),
    );
    vi.stubGlobal('fetch', fetchMock);
    const response = createResponse({ writeResults: [false, true] });
    let settled = false;

    const handledPromise = handleAiProxyRoutes(
      createRequest('POST', { stream: true }),
      response.res,
      new URL('http://127.0.0.1:13210/api/ai-proxy/v1/chat/completions'),
      {} as never,
    ).finally(() => {
      settled = true;
    });
    await flushMicrotasks();

    expect(response.res.flushHeaders).toHaveBeenCalledOnce();
    expect(response.socket.setNoDelay).toHaveBeenCalledWith(true);
    expect(response.body()).toBe('data: one\n\n');
    expect(settled).toBe(false);

    response.res.emit('drain');
    await handledPromise;

    expect(settled).toBe(true);
    expect(response.body()).toBe('data: one\n\ndata: two\n\n');
  });
});