import { Scale, Target, Weight } from 'lucide-react';
import type { LucideIcon } from 'lucide-react';
import { decompose } from '@/lib/demo-kit/reward';
import type { RewardComponent, RewardSpec, RewardValues } from '@/lib/demo-kit/types';
import { Badge } from '@/components/ui/badge';
import { cn } from '@/lib/utils';
import { DASH, formatNumber, formatOrDash } from './format';
import { EditedChip } from './StatStrip';

const ROLE_META: Record<RewardComponent['role'], { label: string; Icon: LucideIcon }> = {
  objective: { label: 'Objective', Icon: Target },
  counterweight: { label: 'Counterweight', Icon: Weight },
  gate: { label: 'Gate', Icon: Scale },
};

export interface RewardBreakdownProps {
  spec: RewardSpec;
  values: RewardValues;
  /** Unweighted diagnostics. Rendered, never summed — the contract is explicit. */
  metrics?: Record<string, number | null>;
  /**
   * Components carrying the weights in force. Pass the output of `reweight()`
   * when the visitor has edited them; omit for the shipped reward.
   */
  components?: readonly RewardComponent[];
  /** Set when `components` came from the visitor rather than the environment. */
  edited?: boolean;
  className?: string;
}

/**
 * `score x weight = value`, per component, plus the total.
 *
 * The counterweight row is called out because it is the row that makes the
 * reward an opinion rather than a scoreboard. Everyone understands "pay for
 * solving it". The interesting engineering is the term that takes points away,
 * and an exec who leaves this page understanding only that has got the point.
 */
export function RewardBreakdown({
  spec,
  values,
  metrics,
  components,
  edited = false,
  className,
}: RewardBreakdownProps) {
  const inForce = components ?? spec.components;
  const { rows, total } = decompose(values, inForce);
  const counterweights = rows.filter((row) => row.role === 'counterweight');

  return (
    <div className={cn('card overflow-hidden', className)}>
      <div className="overflow-x-auto">
        <table className="w-full border-collapse text-sm">
          <caption className="sr-only">
            Reward components, their weights and their contribution to the total score
          </caption>
          <thead>
            <tr className="border-b border-border text-left text-xs uppercase tracking-wide text-muted">
              <th scope="col" className="px-3 py-2 font-medium">
                Component
              </th>
              <th scope="col" className="px-2 py-2 text-right font-medium">
                Score
              </th>
              <th scope="col" className="px-2 py-2 text-right font-medium">
                Weight
              </th>
              <th scope="col" className="px-3 py-2 text-right font-medium">
                Value
              </th>
            </tr>
          </thead>
          <tbody className="divide-y divide-border">
            {rows.map((row) => {
              const role = ROLE_META[row.role];
              const isCounterweight = row.role === 'counterweight';
              return (
                <tr key={row.key} className={cn(isCounterweight && 'bg-accent-subtle/40')}>
                  <th scope="row" className="max-w-[22rem] px-3 py-2.5 text-left font-normal">
                    <span className="flex flex-wrap items-center gap-1.5">
                      <span className="font-medium text-fg">{row.label}</span>
                      <Badge variant={isCounterweight ? 'default' : 'muted'}>
                        <role.Icon className="size-3" aria-hidden="true" />
                        {role.label}
                      </Badge>
                    </span>
                    <span className="mt-0.5 block text-xs leading-snug text-muted">
                      {row.description}
                    </span>
                    <span className="mt-0.5 block font-mono text-[11px] text-muted">{row.key}</span>
                  </th>
                  <td className="nums px-2 py-2.5 text-right align-top font-mono">
                    {row.score === null ? (
                      <span className="text-muted" title="The environment did not score this run">
                        not scored
                      </span>
                    ) : (
                      formatNumber(row.score)
                    )}
                  </td>
                  <td className="nums px-2 py-2.5 text-right align-top font-mono">
                    <span className={cn(edited && 'font-semibold text-accent-fg')}>
                      {formatNumber(row.weight, 2)}
                    </span>
                  </td>
                  <td className="nums px-3 py-2.5 text-right align-top font-mono font-semibold">
                    {row.value === null ? DASH : formatNumber(row.value)}
                  </td>
                </tr>
              );
            })}
          </tbody>
          <tfoot>
            <tr className="border-t-2 border-border bg-surface-2">
              <th scope="row" className="px-3 py-2.5 text-left">
                <span className="flex items-center gap-2 font-semibold">
                  Total reward
                  {edited ? <EditedChip /> : null}
                </span>
              </th>
              <td colSpan={2} />
              <td className="nums px-3 py-2.5 text-right font-mono text-base font-semibold">
                {total === null ? (
                  <span className="text-muted">not scored</span>
                ) : (
                  formatOrDash(total)
                )}
              </td>
            </tr>
          </tfoot>
        </table>
      </div>

      {counterweights.length > 0 ? (
        <p className="border-t border-border bg-accent-subtle/40 px-3 py-2.5 text-xs leading-relaxed text-fg">
          <span className="font-semibold">
            {counterweights.map((row) => row.label).join(' and ')}
          </span>{' '}
          {counterweights.length > 1 ? 'are counterweights' : 'is the counterweight'}: without a
          term pulling the other way, the cheapest way to maximise the objective is a behaviour
          you would never ship, and the model will find it.
        </p>
      ) : null}

      {spec.metrics && spec.metrics.length > 0 ? (
        <div className="border-t border-border px-3 py-2.5">
          <h4 className="text-xs font-semibold uppercase tracking-wide text-muted">
            Diagnostics — reported, never summed
          </h4>
          <dl className="mt-1.5 grid grid-cols-1 gap-x-4 gap-y-1 sm:grid-cols-2">
            {spec.metrics.map((metric) => (
              <div key={metric.key} className="flex items-baseline justify-between gap-2">
                <dt className="text-xs text-muted" title={metric.description}>
                  {metric.label}
                </dt>
                <dd className="nums font-mono text-xs">
                  {formatOrDash(metrics?.[metric.key] ?? null)}
                </dd>
              </div>
            ))}
          </dl>
        </div>
      ) : null}
    </div>
  );
}
