import { useMemo } from 'react';
import {
  BarChart,
  Bar,
  XAxis,
  YAxis,
  Tooltip,
  ResponsiveContainer,
  Cell,
} from 'recharts';

// ---------------------------------------------------------------------------
// Types
// ---------------------------------------------------------------------------

export interface Span {
  spanId: string;
  parentSpanId?: string;
  operationName: string;
  serviceName: string;
  startTime: number; // ms offset from trace start
  duration: number; // ms
  status: string;
  attributes?: Record<string, unknown>;
  events?: Array<{
    name: string;
    timestamp: number;
    attributes?: Record<string, unknown>;
  }>;
}

export interface TraceWaterfallProps {
  spans: Span[];
  onSpanClick?: (span: Span) => void;
  selectedSpanId?: string;
}

// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------

interface FlatSpan extends Span {
  depth: number;
}

function buildTree(spans: Span[]): FlatSpan[] {
  const byId = new Map<string, Span>();
  spans.forEach((s) => byId.set(s.spanId, s));

  const childrenOf = new Map<string, Span[]>();
  const roots: Span[] = [];

  spans.forEach((s) => {
    if (s.parentSpanId && byId.has(s.parentSpanId)) {
      const list = childrenOf.get(s.parentSpanId) ?? [];
      list.push(s);
      childrenOf.set(s.parentSpanId, list);
    } else {
      roots.push(s);
    }
  });

  // Sort roots and children by startTime
  roots.sort((a, b) => a.startTime - b.startTime);
  childrenOf.forEach((list) => list.sort((a, b) => a.startTime - b.startTime));

  const result: FlatSpan[] = [];
  function walk(span: Span, depth: number) {
    result.push({ ...span, depth });
    const kids = childrenOf.get(span.spanId);
    if (kids) kids.forEach((c) => walk(c, depth + 1));
  }
  roots.forEach((r) => walk(r, 0));
  return result;
}

// ---------------------------------------------------------------------------
// Custom bar shape — offset horizontally by startTime
// ---------------------------------------------------------------------------

interface BarShapeProps {
  x?: number;
  y?: number;
  width?: number;
  height?: number;
  payload?: { startTime: number; duration: number; totalDuration: number };
  fill?: string;
}

function OffsetBar(props: BarShapeProps) {
  const { y = 0, height = 0, payload, fill = '#3b82f6' } = props;
  if (!payload) return null;

  const totalDuration = payload.totalDuration || 1;
  // Chart area width — we use the bar's base x and width to derive it.
  // Recharts passes x as the bar origin (0) and width based on the value.
  // Because we set the data value to totalDuration for every bar, width = full chart area.
  const chartWidth = props.width ?? 400;

  const barX = (payload.startTime / totalDuration) * chartWidth;
  const barW = Math.max((payload.duration / totalDuration) * chartWidth, 2);

  return (
    <rect
      x={barX}
      y={y}
      width={barW}
      height={height}
      fill={fill}
      rx={2}
      ry={2}
    />
  );
}

// ---------------------------------------------------------------------------
// Component
// ---------------------------------------------------------------------------

export function TraceWaterfall({ spans, onSpanClick, selectedSpanId }: TraceWaterfallProps) {
  const flatSpans = useMemo(() => buildTree(spans), [spans]);
  const totalDuration = useMemo(() => {
    if (spans.length === 0) return 1;
    return Math.max(...spans.map((s) => s.startTime + s.duration), 1);
  }, [spans]);

  const data = useMemo(
    () =>
      flatSpans.map((s) => ({
        name: `${'  '.repeat(s.depth)}${s.operationName}`,
        value: totalDuration, // full width — custom shape handles offset
        startTime: s.startTime,
        duration: s.duration,
        totalDuration,
        spanId: s.spanId,
        status: s.status,
        depth: s.depth,
      })),
    [flatSpans, totalDuration],
  );

  const chartHeight = Math.max(flatSpans.length * 32 + 40, 120);

  return (
    <ResponsiveContainer width="100%" height={chartHeight}>
      <BarChart
        layout="vertical"
        data={data}
        margin={{ top: 8, right: 16, bottom: 8, left: 0 }}
        barCategoryGap={4}
        onClick={(e) => {
          if (e?.activePayload?.[0]?.payload && onSpanClick) {
            const d = e.activePayload[0].payload as { spanId: string };
            const span = flatSpans.find((s) => s.spanId === d.spanId);
            if (span) onSpanClick(span);
          }
        }}
      >
        <XAxis
          type="number"
          domain={[0, totalDuration]}
          tickFormatter={(v: number) => `${v.toFixed(0)}ms`}
          tick={{ fill: '#94a3b8', fontSize: 11 }}
          axisLine={{ stroke: '#334155' }}
          tickLine={{ stroke: '#334155' }}
        />
        <YAxis
          type="category"
          dataKey="name"
          width={220}
          tick={{ fill: '#94a3b8', fontSize: 11 }}
          axisLine={false}
          tickLine={false}
        />
        <Tooltip
          cursor={{ fill: 'rgba(59,130,246,0.08)' }}
          contentStyle={{
            backgroundColor: '#1e293b',
            border: '1px solid #334155',
            borderRadius: 6,
            fontSize: 12,
            color: '#e2e8f0',
          }}
          formatter={(_v: unknown, _n: unknown, props: { payload?: { duration: number; startTime: number } }) => {
            const p = props.payload;
            if (!p) return ['', ''];
            return [`${p.duration.toFixed(1)}ms (offset ${p.startTime.toFixed(1)}ms)`, 'Duration'];
          }}
        />
        <Bar dataKey="value" shape={<OffsetBar />} isAnimationActive={false}>
          {data.map((d) => {
            const isError = d.status === 'error' || d.status === 'ERROR';
            const isSelected = d.spanId === selectedSpanId;
            let fill = isError ? '#ef4444' : '#3b82f6';
            if (isSelected) fill = isError ? '#fca5a5' : '#60a5fa';
            return <Cell key={d.spanId} fill={fill} />;
          })}
        </Bar>
      </BarChart>
    </ResponsiveContainer>
  );
}
