import type { componentsV1 } from "@arizeai/phoenix-client";
import { describe, expect, it } from "vitest";

import {
  type ExperimentWithRuns,
  formatExperimentJsonOutput,
  formatExperimentOutput,
} from "../src/commands/formatExperiment";
import { formatExperimentsOutput } from "../src/commands/formatExperiments";

type Experiment = componentsV1["schemas"]["Experiment"];
type ExperimentRun = componentsV1["schemas"]["ExperimentRun"];

// Mock experiment data
const mockExperiment1: Experiment = {
  id: "exp-abc123",
  dataset_id: "ds-123",
  dataset_version_id: "v1-abc",
  repetitions: 1,
  metadata: {},
  project_name: "my-project",
  created_at: "2026-01-15T10:00:00.000Z",
  updated_at: "2026-01-15T12:00:00.000Z",
  example_count: 10,
  successful_run_count: 8,
  failed_run_count: 2,
  missing_run_count: 0,
};

const mockExperiment2: Experiment = {
  id: "exp-xyz789",
  dataset_id: "ds-456",
  dataset_version_id: "v2-xyz",
  repetitions: 3,
  metadata: { model: "gpt-4" },
  project_name: null,
  created_at: "2026-01-14T08:00:00.000Z",
  updated_at: "2026-01-14T09:00:00.000Z",
  example_count: 5,
  successful_run_count: 15,
  failed_run_count: 0,
  missing_run_count: 0,
};

const mockRun1: ExperimentRun = {
  id: "run-001",
  experiment_id: "exp-abc123",
  dataset_example_id: "example-1",
  output: { result: "success" },
  repetition_number: 1,
  start_time: "2026-01-15T10:00:00.000Z",
  end_time: "2026-01-15T10:00:01.500Z",
  trace_id: "trace-abc",
  error: null,
};

const mockRunWithError: ExperimentRun = {
  id: "run-002",
  experiment_id: "exp-abc123",
  dataset_example_id: "example-2",
  output: null,
  repetition_number: 1,
  start_time: "2026-01-15T10:00:02.000Z",
  end_time: "2026-01-15T10:00:02.100Z",
  trace_id: null,
  error: "ValueError: Invalid input",
};

describe("Experiments Formatting", () => {
  describe("formatExperimentsOutput - raw", () => {
    it("should format as compact JSON", () => {
      const output = formatExperimentsOutput({
        experiments: [mockExperiment1],
        format: "raw",
      });

      expect(output).toBe(JSON.stringify([mockExperiment1]));
      expect(output).not.toContain("\n");
    });

    it("should handle empty array", () => {
      const output = formatExperimentsOutput({
        experiments: [],
        format: "raw",
      });

      expect(output).toBe("[]");
    });
  });

  describe("formatExperimentsOutput - json", () => {
    it("should format as pretty JSON", () => {
      const output = formatExperimentsOutput({
        experiments: [mockExperiment1],
        format: "json",
      });

      expect(output).toBe(JSON.stringify([mockExperiment1], null, 2));
      expect(output).toContain("\n");
    });

    it("should handle multiple experiments", () => {
      const output = formatExperimentsOutput({
        experiments: [mockExperiment1, mockExperiment2],
        format: "json",
      });

      expect(output).toBe(
        JSON.stringify([mockExperiment1, mockExperiment2], null, 2)
      );
    });
  });

  describe("formatExperimentsOutput - pretty", () => {
    it("should format experiments as a table with column headers", () => {
      const output = formatExperimentsOutput({
        experiments: [mockExperiment1],
        format: "pretty",
      });

      expect(output).toContain("id");
      expect(output).toContain("exp-abc123");
      expect(output).toContain("my-project");
      expect(output).toContain("ds-123");
      expect(output).toContain("10"); // examples
      expect(output).toContain("8"); // successful
      expect(output).toContain("2"); // failed
      // table borders
      expect(output).toContain("┌");
      expect(output).toContain("┘");
    });

    it("should handle experiments without project name", () => {
      const output = formatExperimentsOutput({
        experiments: [mockExperiment2],
        format: "pretty",
      });

      expect(output).toContain("exp-xyz789");
      expect(output).toContain("ds-456");
    });

    it("should handle empty array", () => {
      const output = formatExperimentsOutput({
        experiments: [],
        format: "pretty",
      });

      expect(output).toBe("No experiments found");
    });

    it("should show multiple experiments", () => {
      const output = formatExperimentsOutput({
        experiments: [mockExperiment1, mockExperiment2],
        format: "pretty",
      });

      expect(output).toContain("exp-abc123");
      expect(output).toContain("exp-xyz789");
    });
  });
});

describe("Experiment (Single) Formatting", () => {
  const mockExperimentWithRuns: ExperimentWithRuns = {
    experiment: mockExperiment1,
    runs: [mockRun1, mockRunWithError],
  };

  describe("formatExperimentOutput - raw", () => {
    it("should format as compact JSON", () => {
      const output = formatExperimentOutput({
        data: mockExperimentWithRuns,
        format: "raw",
      });

      expect(output).toBe(JSON.stringify(mockExperimentWithRuns));
    });
  });

  describe("formatExperimentOutput - json", () => {
    it("should format as pretty JSON", () => {
      const output = formatExperimentOutput({
        data: mockExperimentWithRuns,
        format: "json",
      });

      expect(output).toBe(JSON.stringify(mockExperimentWithRuns, null, 2));
    });
  });

  describe("formatExperimentOutput - pretty", () => {
    it("should format experiment with runs in human-readable format", () => {
      const output = formatExperimentOutput({
        data: mockExperimentWithRuns,
        format: "pretty",
      });

      expect(output).toContain(
        "┌─ Experiment: exp-abc123 [Project: my-project]"
      );
      expect(output).toContain("│  Dataset ID: ds-123");
      expect(output).toContain("│  Run Summary:");
      expect(output).toContain("│    ✓ Successful: 8");
      expect(output).toContain("│  Runs (2):");
      expect(output).toContain("│    ✓ run-001");
      expect(output).toContain("[trace: trace-abc]");
      expect(output).toContain("│    ✗ run-002");
      expect(output).toContain("Error: ValueError: Invalid input");
    });

    it("should show output preview", () => {
      const output = formatExperimentOutput({
        data: mockExperimentWithRuns,
        format: "pretty",
      });

      expect(output).toContain('Output: {"result":"success"}');
    });
  });
});

describe("Experiment JSON Output Formatting", () => {
  const mockJsonData = JSON.stringify([
    {
      id: "run-001",
      example_id: "example-1",
      input: { question: "What is 2+2?" },
      output: { answer: "4" },
      expected_output: { answer: "4" },
      start_time: "2026-01-15T10:00:00.000Z",
      end_time: "2026-01-15T10:00:01.500Z",
      trace_id: "trace-abc",
      evaluations: {
        correctness: { score: 1.0, label: "correct" },
      },
    },
    {
      id: "run-002",
      example_id: "example-2",
      input: { question: "What is 3+3?" },
      output: null,
      error: "Timeout error",
      start_time: "2026-01-15T10:00:02.000Z",
      end_time: "2026-01-15T10:00:05.000Z",
    },
  ]);

  describe("formatExperimentJsonOutput - raw", () => {
    it("should format as compact JSON", () => {
      const output = formatExperimentJsonOutput({
        jsonData: mockJsonData,
        format: "raw",
      });

      const parsed = JSON.parse(output);
      expect(Array.isArray(parsed)).toBe(true);
      expect(parsed).toHaveLength(2);
    });
  });

  describe("formatExperimentJsonOutput - json", () => {
    it("should format as pretty JSON", () => {
      const output = formatExperimentJsonOutput({
        jsonData: mockJsonData,
        format: "json",
      });

      expect(output).toContain("\n");
      expect(output).toContain("  ");
      const parsed = JSON.parse(output);
      expect(parsed).toHaveLength(2);
    });
  });

  describe("formatExperimentJsonOutput - pretty", () => {
    it("should format experiment runs in human-readable format", () => {
      const output = formatExperimentJsonOutput({
        jsonData: mockJsonData,
        format: "pretty",
      });

      expect(output).toContain("Experiment Runs (2):");
      expect(output).toContain("┌─ ✓ Run: run-001");
      expect(output).toContain("│  Example ID: example-1");
      expect(output).toContain("│  Trace ID: trace-abc");
      expect(output).toContain('│  Input: {"question":"What is 2+2?"}');
      expect(output).toContain('│  Output: {"answer":"4"}');
      expect(output).toContain("│  Evaluations:");
      expect(output).toContain('- correctness: score=1, label="correct"');
    });

    it("should show errors for failed runs", () => {
      const output = formatExperimentJsonOutput({
        jsonData: mockJsonData,
        format: "pretty",
      });

      expect(output).toContain("┌─ ✗ Run: run-002");
      expect(output).toContain("│  Error: Timeout error");
    });
  });
});
