/**
 * Copyright 2024 Google LLC
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

import {
  GenerateRequestData,
  GenerateResponseChunkData,
  GenerateResponseData,
  GenerateResponseSchema,
  Part,
} from '@genkit-ai/tools-common';
import {
  BaseRuntimeManager,
  GenkitToolsError,
} from '@genkit-ai/tools-common/manager';
import { findProjectRoot, logger } from '@genkit-ai/tools-common/utils';
import { Command } from 'commander';
import { readFileSync } from 'fs';
import { resolve } from 'path';
import { parse } from 'yaml';
import { startDevProcessManager, startManager } from '../utils/manager-utils';

interface TestOptions {
  supports: string;
  fromFile?: string;
}

type TestCase = {
  name: string;
  input: GenerateRequestData;
  validators: string[];
  stream?: boolean;
};

type TestSuite = {
  model: string;
  supports?: string[];
  tests?: TestCase[];
};

const getMessageText = (response: GenerateResponseData): string | undefined => {
  const message = response.message || response.candidates?.[0]?.message;
  return message?.content?.[0]?.text;
};

const getMessageContent = (response: GenerateResponseData) => {
  const message = response.message || response.candidates?.[0]?.message;
  return message?.content;
};

const getMediaPart = (response: GenerateResponseData) => {
  const content = getMessageContent(response);
  return content?.find((p: Part) => p.media);
};

const imageBase64 =
  ' ';

const VALIDATORS: Record<
  string,
  (
    response: GenerateResponseData,
    arg?: string,
    chunks?: GenerateResponseChunkData[]
  ) => void
> = {
  'has-tool-request': (response, toolName) => {
    const content = getMessageContent(response);
    if (!content || !Array.isArray(content)) {
      throw new Error(
        `Response missing message content. Full response: ${JSON.stringify(
          response,
          null,
          2
        )}`
      );
    }
    const toolRequest = content.find((c: Part) => c.toolRequest);
    if (!toolRequest) {
      throw new Error(
        `Model did not return a tool request. Content: ${JSON.stringify(
          content,
          null,
          2
        )}`
      );
    }
    if (toolName && toolRequest.toolRequest?.name !== toolName) {
      throw new Error(
        `Expected tool request '${toolName}', got '${toolRequest.toolRequest?.name}'`
      );
    }
  },
  'valid-json': (response) => {
    const content = getMessageContent(response);
    if (!content || !Array.isArray(content)) {
      throw new Error(
        `Response missing message content. Full response: ${JSON.stringify(
          response,
          null,
          2
        )}`
      );
    }
    const textPart = content.find((c: Part) => c.text);
    if (!textPart) {
      throw new Error(
        `Model did not return text content for JSON. Content: ${JSON.stringify(
          content,
          null,
          2
        )}`
      );
    }
    try {
      JSON.parse(textPart.text!);
    } catch (e) {
      throw new Error(
        `Response text is not valid JSON. Text: ${textPart.text}`
      );
    }
  },
  'text-includes': (response, expected) => {
    const text = getMessageText(response);
    if (
      !text ||
      (expected && !text.toLowerCase().includes(expected.toLowerCase()))
    ) {
      throw new Error(
        `Response text does not include '${expected}'. Text: ${text}`
      );
    }
  },
  'stream-text-includes': (response, expected, chunks) => {
    if (!chunks || chunks.length === 0) {
      throw new Error('Streaming expected but no chunks were received');
    }

    const streamedText = chunks
      .map((c) => c.content?.find((p: any) => p.text)?.text || '')
      .join('');

    if (expected && !streamedText.includes(expected)) {
      throw new Error(`Streaming response did not include ${expected}'`);
    }
  },
  'stream-has-tool-request': (response, toolName, chunks) => {
    if (!chunks || chunks.length === 0) {
      throw new Error('Streaming expected but no chunks were received');
    }

    const hasTool = chunks.some((c) =>
      c.content?.some((p: any) => !!p.toolRequest)
    );
    if (!hasTool) {
      throw new Error('No tool request found in the streamed chunks');
    }

    if (toolName) {
      VALIDATORS['has-tool-request'](response, toolName);
    }
  },
  'stream-valid-json': (response, arg, chunks) => {
    if (!chunks || chunks.length === 0) {
      throw new Error('Streaming expected but no chunks were received');
    }

    const streamedText = chunks
      .map((c) => c.content?.find((p: any) => p.text)?.text || '')
      .join('');

    if (!streamedText.trim()) {
      throw new Error('Streamed response contained no text');
    }
    try {
      JSON.parse(streamedText);
    } catch (e) {
      throw new Error(`Streamed text is not valid JSON: ${streamedText}`);
    }
  },
  'text-starts-with': (response, expected) => {
    const text = getMessageText(response);
    if (!text || (expected && !text.trim().startsWith(expected))) {
      throw new Error(
        `Response text does not start with '${expected}'. Text: ${text}`
      );
    }
  },
  'text-not-empty': (response) => {
    const text = getMessageText(response);
    if (!text || text.trim().length === 0) {
      throw new Error('Response text is empty');
    }
  },
  'valid-media': (response, type) => {
    const mediaPart = getMediaPart(response);
    if (!mediaPart) {
      throw new Error(`Model did not return ${type || 'media'} part.`);
    }
    if (type) {
      if (
        mediaPart.media?.contentType &&
        !mediaPart.media.contentType.startsWith(`${type}/`)
      ) {
        throw new Error(
          `Expected ${type} content type, got ${mediaPart.media.contentType}`
        );
      }
    }
    if (type === 'image') {
      const url = mediaPart.media?.url;
      if (!url) throw new Error('Media part missing URL');
      if (url.startsWith('data:')) {
        if (!url.startsWith('data:image/')) {
          throw new Error('Invalid data URL content type for image');
        }
      } else if (url.startsWith('http')) {
        try {
          new URL(url);
        } catch (e) {
          throw new Error(`Invalid URL: ${url}`);
        }
      } else {
        throw new Error(`Unknown URL format: ${url}`);
      }
    }
  },
  reasoning: (response) => {
    const content = getMessageContent(response);

    if (!content || !Array.isArray(content)) {
      throw new Error(`Response is missing message content`);
    }

    const hasReasoning = content.some((p: any) => !!p.reasoning);
    if (!hasReasoning) {
      throw new Error(`reasoning content not found`);
    }
  },
};

const TEST_CASES: Record<string, TestCase> = {
  'tool-request': {
    name: 'Tool Request Conformance',
    input: {
      messages: [
        {
          role: 'user',
          content: [{ text: 'What is the weather in New York? Use the tool.' }],
        },
      ],
      tools: [
        {
          name: 'weather',
          description: 'Get the weather for a city',
          inputSchema: {
            type: 'object',
            properties: {
              city: { type: 'string' },
            },
            required: ['city'],
          },
        },
      ],
    },
    validators: ['has-tool-request:weather'],
  },
  'structured-output': {
    name: 'Structured Output Conformance',
    input: {
      messages: [
        {
          role: 'user',
          content: [{ text: 'Generate a profile for John Doe.' }],
        },
      ],
      output: {
        format: 'json',
        schema: {
          type: 'object',
          properties: {
            name: { type: 'string' },
            age: { type: 'number' },
          },
          required: ['name', 'age'],
        },
        constrained: true,
      },
    },
    validators: ['valid-json'],
  },
  multiturn: {
    name: 'Multiturn Conformance',
    input: {
      messages: [
        { role: 'user', content: [{ text: 'My name is Genkit.' }] },
        { role: 'model', content: [{ text: 'Hello Genkit.' }] },
        { role: 'user', content: [{ text: 'What is my name?' }] },
      ],
    },
    validators: ['text-includes:Genkit'],
  },
  'streaming-multiturn': {
    name: 'Multiturn Conformance with streaming',
    stream: true,
    input: {
      messages: [
        { role: 'user', content: [{ text: 'My name is Genkit.' }] },
        { role: 'model', content: [{ text: 'Hello Genkit.' }] },
        { role: 'user', content: [{ text: 'What is my name?' }] },
      ],
    },
    validators: ['stream-text-includes:Genkit'],
  },
  'streaming-tool-request': {
    name: 'Tool Request Conformance with streaming',
    stream: true,
    input: {
      messages: [
        {
          role: 'user',
          content: [
            { text: 'What is the weather in New York? Use the weather tool' },
          ],
        },
      ],
      tools: [
        {
          name: 'weather',
          description: 'Get the weather for a city',
          inputSchema: {
            type: 'object',
            properties: {
              city: { type: 'string' },
            },
            required: ['city'],
          },
        },
      ],
    },
    validators: ['stream-has-tool-request:weather'],
  },
  'streaming-structured-output': {
    name: 'Structured Output Conformance with streaming',
    stream: true,
    input: {
      messages: [
        {
          role: 'user',
          content: [{ text: 'Generate a movie review for John Wick' }],
        },
      ],
      output: {
        format: 'json',
        schema: {
          type: 'object',
          properties: {
            name: { type: 'string' },
            rating: { type: 'number' },
          },
          required: ['name', 'rating'],
        },
        constrained: true,
      },
    },
    validators: ['stream-valid-json'],
  },
  'system-role': {
    name: 'System Role Conformance',
    input: {
      messages: [
        {
          role: 'system',
          content: [
            {
              text: "IMPORTANT: your response are machine processed, always start/prefix your response with 'RESPONSE:', ex: 'RESPONSE: hello'",
            },
          ],
        },
        { role: 'user', content: [{ text: 'hello' }] },
      ],
    },
    validators: ['text-starts-with:RESPONSE:'],
  },
  'input-image-base64': {
    name: 'Image Input (Base64) Conformance',
    input: {
      messages: [
        {
          role: 'user',
          content: [
            { text: 'What text do you see in this image?' },
            {
              media: {
                url: `data:image/png;base64,${imageBase64}`,
                contentType: 'image/png',
              },
            },
          ],
        },
      ],
    },
    validators: ['text-includes:genkit'],
  },
  'input-image-url': {
    name: 'Image Input (URL) Conformance',
    input: {
      messages: [
        {
          role: 'user',
          content: [
            { text: 'What is this logo?' },
            {
              media: {
                url: 'https://www.google.com/images/branding/googlelogo/2x/googlelogo_color_272x92dp.png',
                contentType: 'image/png',
              },
            },
          ],
        },
      ],
    },
    validators: ['text-includes:google'],
  },
  'input-video-youtube': {
    name: 'Video Input (YouTube) Conformance',
    input: {
      messages: [
        {
          role: 'user',
          content: [
            { text: 'Describe this video.' },
            {
              media: {
                url: 'https://www.youtube.com/watch?v=3p1P5grjXIQ',
                contentType: 'video/mp4',
              },
            },
          ],
        },
      ],
    },
    validators: ['text-not-empty'],
  },
  'output-audio': {
    name: 'Audio Output (TTS) Conformance',
    input: {
      messages: [{ role: 'user', content: [{ text: 'Say hello.' }] }],
    },
    validators: ['valid-media:audio'],
  },
  'output-image': {
    name: 'Image Output (Generation) Conformance',
    input: {
      messages: [
        {
          role: 'user',
          content: [{ text: 'Generate an image of a cat.' }],
        },
      ],
    },
    validators: ['valid-media:image'],
  },
};

async function waitForRuntime(manager: BaseRuntimeManager) {
  // Poll for runtimes
  for (let i = 0; i < 20; i++) {
    if (manager.listRuntimes().length > 0) return;
    await new Promise((r) => setTimeout(r, 500));
  }
  logger.warn('Runtime not detected after 10 seconds.');
}

/**
 * Waits until all model actions referenced by the test suites are registered
 * with the runtime's reflection server. This prevents 404 errors caused by
 * dispatching tests before the runtime has finished registering actions.
 */
async function waitForActions(
  manager: BaseRuntimeManager,
  suites: TestSuite[]
): Promise<void> {
  const requiredKeys = new Set<string>();
  for (const suite of suites) {
    if (suite.model) {
      const key = suite.model.startsWith('/')
        ? suite.model
        : `/model/${suite.model}`;
      requiredKeys.add(key);
    }
  }
  if (requiredKeys.size === 0) return;

  const maxAttempts = 60; // 30 seconds total
  const delayMs = 500;

  for (let attempt = 0; attempt < maxAttempts; attempt++) {
    try {
      const actions = await manager.listActions();
      const registeredKeys = new Set(Object.keys(actions));
      const missing = [...requiredKeys].filter((k) => !registeredKeys.has(k));
      if (missing.length === 0) {
        logger.info(
          `All ${requiredKeys.size} model actions registered. Starting tests.`
        );
        return;
      }
      if (attempt % 10 === 0 && attempt > 0) {
        logger.info(
          `Waiting for ${missing.length} model action(s) to register: ${missing.join(', ')}...`
        );
      }
    } catch (e) {
      // Actions endpoint not ready yet, keep polling.
      logger.debug(`Polling for actions failed, will retry: ${e}`);
    }
    await new Promise((r) => setTimeout(r, delayMs));
  }
  logger.warn(
    'Not all model actions registered after 30 seconds. Proceeding anyway.'
  );
}

async function runTest(
  manager: BaseRuntimeManager,
  model: string,
  testCase: TestCase
): Promise<boolean> {
  logger.info(`Running test: ${testCase.name}...`);
  try {
    // Adjust model name if needed (e.g. /model/ prefix)
    const modelKey = model.startsWith('/') ? model : `/model/${model}`;
    const shouldStream = !!testCase.stream;

    const collectedChunks: any[] = [];

    const actionResponse = await manager.runAction(
      {
        key: modelKey,
        input: testCase.input,
      },
      shouldStream
        ? (chunk) => {
            collectedChunks.push(chunk);
          }
        : undefined
    );

    if (shouldStream && collectedChunks.length === 0) {
      throw new Error('Streaming requested but no chunks received.');
    }
    const response = GenerateResponseSchema.parse(actionResponse.result);

    for (const v of testCase.validators) {
      const [valName, ...args] = v.split(':');
      const arg = args.join(':');
      const validator = VALIDATORS[valName];
      if (!validator) throw new Error(`Unknown validator: ${valName}`);
      validator(response, arg, collectedChunks);
    }

    logger.info(`✅ Passed: ${testCase.name}`);
    return true;
  } catch (e) {
    if (e instanceof GenkitToolsError) {
      logger.error(
        `❌ Failed: ${testCase.name} - ${
          e.data?.stack || JSON.stringify(e.data?.details) || e
        }`
      );
    } else if (e instanceof Error) {
      logger.error(`❌ Failed: ${testCase.name} - ${e.message}`);
    } else {
      logger.error(`❌ Failed: ${testCase.name} - ${JSON.stringify(e)}`);
    }
    return false;
  }
}

async function runTestSuite(
  manager: BaseRuntimeManager,
  suite: TestSuite,
  defaultSupports: string[]
): Promise<{ passed: number; failed: number }> {
  const supports = suite.supports || (suite.tests ? [] : defaultSupports);

  logger.info(`Testing model: ${suite.model}`);

  const promises: Promise<boolean>[] = [];

  // Built-in conformance tests
  for (const support of supports) {
    const testCase = TEST_CASES[support];
    if (testCase) {
      promises.push(runTest(manager, suite.model, testCase));
    } else {
      logger.warn(`Unknown capability: ${support}`);
    }
  }

  // Custom tests
  if (suite.tests) {
    for (const test of suite.tests) {
      const customTestCase: TestCase = {
        name: test.name || 'Custom Test',
        input: test.input,
        validators: test.validators || [],
        stream: test.stream,
      };
      promises.push(runTest(manager, suite.model, customTestCase));
    }
  }

  const results = await Promise.all(promises);
  const passed = results.filter((r) => r).length;
  const failed = results.filter((r) => !r).length;

  return { passed, failed };
}

export const devTestModel = new Command('dev:test-model')
  .description('Test a model against the Genkit model specification')
  .argument('[modelOrCmd]', 'Model name or command')
  .argument('[args...]', 'Command arguments')
  .option(
    '--supports <list>',
    'Comma-separated list of supported capabilities (tool-request, structured-output, multiturn, system-role, input-image-base64, input-image-url, input-video-youtube, output-audio, output-image, streaming-multiturn, reasoning)',
    'tool-request,structured-output,multiturn,system-role,input-image-base64,input-image-url,streaming-multiturn,streaming-tool-request,streaming-structured-output'
  )
  .option('--from-file <file>', 'Path to a file containing test payloads')
  .action(
    async (
      modelOrCmd: string | undefined,
      args: string[] | undefined,
      options: TestOptions
    ) => {
      const projectRoot = await findProjectRoot();

      let cmd: string[] = [];
      let defaultModelName: string | undefined;

      if (options.fromFile) {
        if (modelOrCmd) cmd.push(modelOrCmd);
        if (args) cmd.push(...args);
      } else {
        if (!modelOrCmd) {
          logger.error('Model name is required unless --from-file is used.');
          process.exitCode = 1;
          return;
        }
        defaultModelName = modelOrCmd;
        if (args) cmd = args;
      }

      let manager: BaseRuntimeManager;

      if (cmd.length > 0) {
        const result = await startDevProcessManager(
          projectRoot,
          cmd[0],
          cmd.slice(1)
        );
        manager = result.manager;
      } else {
        manager = await startManager({ projectRoot, manageHealth: false });
      }

      await waitForRuntime(manager);

      try {
        let totalPassed = 0;
        let totalFailed = 0;

        let suites: TestSuite[] = [];

        if (options.fromFile) {
          const filePath = resolve(projectRoot, options.fromFile);
          const fileContent = readFileSync(filePath, 'utf-8');
          let parsed;
          if (filePath.endsWith('.yaml') || filePath.endsWith('.yml')) {
            parsed = parse(fileContent);
          } else {
            parsed = JSON.parse(fileContent);
          }
          suites = Array.isArray(parsed) ? parsed : [parsed];
        } else {
          if (!defaultModelName) throw new Error('Model name required');
          suites = [{ model: defaultModelName }];
        }

        const defaultSupports = options.supports
          .split(',')
          .map((s) => s.trim());

        // Wait for all model actions to be registered before dispatching tests.
        // This prevents 404 errors when the runtime is still initializing.
        await waitForActions(manager, suites);

        for (const suite of suites) {
          if (!suite.model) {
            logger.error('Model name required in test suite.');
            totalFailed++;
            continue;
          }
          const { passed, failed } = await runTestSuite(
            manager,
            suite,
            defaultSupports
          );
          totalPassed += passed;
          totalFailed += failed;
        }

        logger.info('--------------------------------------------------');
        logger.info(
          `Tests Completed: ${totalPassed} Passed, ${totalFailed} Failed`
        );

        if (totalFailed > 0) {
          process.exitCode = 1;
        }
      } catch (e) {
        logger.error('Error running tests:', e);
        process.exitCode = 1;
      } finally {
        if (manager) {
          await manager.stop();
        }
      }
    }
  );
