import { Bar, BarChart, CartesianGrid, LabelList, XAxis, YAxis } from "recharts";
import { Card, CardContent } from "@/components/ui/card";
import type { ChartConfig } from "@/components/ui/chart";
import {
  ChartContainer,
  ChartTooltip,
  ChartTooltipContent,
} from "@/components/ui/chart";

export interface BenchmarkRow {
  compileGroup: string;
  eagerPrefill: number;
  eagerDecode: number;
  eagerTotal: number;
  groupedMmPrefill: number;
  groupedMmDecode: number;
  groupedMmTotal: number;
  batchedMmPrefill: number;
  batchedMmDecode: number;
  batchedMmTotal: number;
  groupedBatchedPrefill: number;
  groupedBatchedDecode: number;
  groupedBatchedTotal: number;
}

interface BenchmarkChartProps {
  data: BenchmarkRow[];
}

const chartConfig = {
  eagerPrefill: {
    label: "eager prefill",
    color: "#2563eb",
  },
  eagerDecode: {
    label: "eager decode",
    color: "#93c5fd",
  },
  groupedMmPrefill: {
    label: "grouped_mm prefill",
    color: "#16a34a",
  },
  groupedMmDecode: {
    label: "grouped_mm decode",
    color: "#86efac",
  },
  batchedMmPrefill: {
    label: "batched_mm prefill",
    color: "#7c3aed",
  },
  batchedMmDecode: {
    label: "batched_mm decode",
    color: "#d8b4fe",
  },
  groupedBatchedPrefill: {
    label: "grouped+batched prefill",
    color: "#ea580c",
  },
  groupedBatchedDecode: {
    label: "grouped+batched decode",
    color: "#fdba74",
  },
} satisfies ChartConfig;

export function BenchmarkChart({
  data,
}: BenchmarkChartProps) {
  return (
    <Card className="border-0 shadow-none gap-0">
      <CardContent className="px-0 pr-6 pt-0">
        <ChartContainer config={chartConfig} style={{ height: 350, width: '100%' }}>
          <BarChart accessibilityLayer data={data}>
            <CartesianGrid vertical={false} />
            <XAxis
              dataKey="compileGroup"
              tickLine={false}
              tickMargin={10}
              axisLine={false}
            />
            <YAxis
              tickLine={false}
              axisLine={false}
              width={90}
              tickFormatter={(value) => `${value}ms`}
              label={{
                value: "generation latency (ms)",
                angle: -90,
                position: "left",
                offset: 70,
                dx: -28,
              }}
            />
            <ChartTooltip
              cursor={false}
              shared={false}
              isAnimationActive={false}
              content={<ChartTooltipContent indicator="dashed" />}
            />
            <Bar
              dataKey="eagerPrefill"
              fill="var(--color-eagerPrefill)"
              stackId="eager"
              radius={[0, 0, 4, 4]}
            />
            <Bar
              dataKey="eagerDecode"
              fill="var(--color-eagerDecode)"
              stackId="eager"
              radius={[4, 4, 0, 0]}
            >
              <LabelList
                dataKey="eagerTotal"
                position="top"
                formatter={(value) =>
                  value == null ? "" : Math.round(Number(value)).toString()
                }
              />
            </Bar>
            <Bar
              dataKey="groupedMmPrefill"
              fill="var(--color-groupedMmPrefill)"
              stackId="groupedMm"
              radius={[0, 0, 4, 4]}
            />
            <Bar
              dataKey="groupedMmDecode"
              fill="var(--color-groupedMmDecode)"
              stackId="groupedMm"
              radius={[4, 4, 0, 0]}
            >
              <LabelList
                dataKey="groupedMmTotal"
                position="top"
                formatter={(value) =>
                  value == null ? "" : Math.round(Number(value)).toString()
                }
              />
            </Bar>
            <Bar
              dataKey="batchedMmPrefill"
              fill="var(--color-batchedMmPrefill)"
              stackId="batchedMm"
              radius={[0, 0, 4, 4]}
            />
            <Bar
              dataKey="batchedMmDecode"
              fill="var(--color-batchedMmDecode)"
              stackId="batchedMm"
              radius={[4, 4, 0, 0]}
            >
              <LabelList
                dataKey="batchedMmTotal"
                position="top"
                formatter={(value) =>
                  value == null ? "" : Math.round(Number(value)).toString()
                }
              />
            </Bar>
            <Bar
              dataKey="groupedBatchedPrefill"
              fill="var(--color-groupedBatchedPrefill)"
              stackId="groupedBatched"
              radius={[0, 0, 4, 4]}
            />
            <Bar
              dataKey="groupedBatchedDecode"
              fill="var(--color-groupedBatchedDecode)"
              stackId="groupedBatched"
              radius={[4, 4, 0, 0]}
            >
              <LabelList
                dataKey="groupedBatchedTotal"
                position="top"
                formatter={(value) =>
                  value == null ? "" : Math.round(Number(value)).toString()
                }
              />
            </Bar>
          </BarChart>
        </ChartContainer>
      </CardContent>
    </Card>
  );
}