offscale/web/PolicyAssistant.tsx

import { useMemo, useState } from 'react';
import { Check, Loader2, Sparkles, Wand2, X } from 'lucide-react';
import { useHeadscalePolicyAssist, assistFailure } from './useHeadscalePolicy';
import { collapseUnchanged, diffCounts, diffLines } from './diff';
import { Button, Card, ErrorNote } from './Cards';

// Ask for a policy change in English; read the diff; decide.
//
// The whole point of this panel is the middle step. The model is good at the grammar — HuJSON, tagOwners,
// the src/dst shapes — and has no idea which of the owner's machines matter, so its proposal is a draft to
// be read, not an answer to be trusted. Nothing here writes to Headscale: Apply puts the text in the editor
// above and the existing Save button is still the only thing that leaves the browser.
//
// The diff is against what is CURRENTLY in the editor, which is also what was sent up, so it always shows
// exactly what accepting would change on screen — including edits the owner made and hasn't saved.

const EXAMPLES = [
  'let everyone reach the machines tagged tag:server on port 22',
  'stop the phones from reaching anything except the DNS server',
  'add a group for family with just my own user in it',
];

const DiffBody = ({ before, after }: { before: string; after: string }) => {
  const lines = useMemo(() => diffLines(before, after), [before, after]);
  const rows = useMemo(() => collapseUnchanged(lines), [lines]);
  const { added, removed } = useMemo(() => diffCounts(lines), [lines]);

  if (!added && !removed) {
    return <p className="px-3 py-2.5 text-[11px] text-zinc-500">No change — the proposal matches what you have.</p>;
  }

  return (
    <>
      <div className="flex items-center gap-3 border-b border-white/10 px-3 py-1.5 text-[11px]">
        <span className="text-emerald-400">+{added}</span>
        <span className="text-red-400">−{removed}</span>
        <span className="text-zinc-600">unchanged lines collapsed</span>
      </div>
      <div className="max-h-72 overflow-auto p-1 font-mono text-[11px] leading-relaxed">
        {rows.map((row, index) =>
          row === null ? (
            <div key={index} className="px-2 py-1 text-center text-zinc-700 select-none">
              ⋯
            </div>
          ) : (
            <div
              key={index}
              className={`px-2 whitespace-pre-wrap ${
                row.kind === 'add'
                  ? 'bg-emerald-500/10 text-emerald-300'
                  : row.kind === 'remove'
                    ? 'bg-red-500/10 text-red-300'
                    : 'text-zinc-500'
              }`}
            >
              {row.kind === 'add' ? '+' : row.kind === 'remove' ? '−' : ' '} {row.text}
            </div>
          ),
        )}
      </div>
    </>
  );
};

type PolicyAssistantProps = {
  /** The text on screen right now. Sent up as the base, and diffed against. */
  policy: string;
  /** Accepting a proposal — puts it in the editor's draft. Never saves. */
  onApply: (policy: string) => void;
  disabled?: boolean;
};

export const PolicyAssistant = ({ policy, onApply, disabled }: PolicyAssistantProps) => {
  const assist = useHeadscalePolicyAssist();
  const [prompt, setPrompt] = useState('');
  // The proposal lives here, not in the mutation's `data`, so that discarding it is a real state change and
  // a second ask doesn't briefly show the previous answer against the new base.
  const [proposal, setProposal] = useState<{ explanation: string; policy: string } | null>(null);

  const ask = async () => {
    const request = prompt.trim();
    if (!request || assist.isPending) return;
    setProposal(null);
    try {
      setProposal(await assist.mutateAsync({ prompt: request, policy }));
    } catch {
      // Rendered from `assist.error` below — mutateAsync rejecting is the same failure twice.
    }
  };

  const apply = () => {
    if (!proposal) return;
    onApply(proposal.policy);
    setProposal(null);
    setPrompt('');
    assist.reset();
  };

  return (
    <Card>
      <div className="flex items-center gap-2 border-b border-white/10 px-3 py-2">
        <Sparkles className="h-3.5 w-3.5 text-primary" />
        <span className="text-xs font-medium text-zinc-200">Describe the change</span>
        <span className="ml-auto text-[11px] text-zinc-600">Proposes a document — never saves it</span>
      </div>

      <div className="flex flex-col gap-2 p-3">
        <textarea
          value={prompt}
          onChange={(ev) => setPrompt(ev.target.value)}
          onKeyDown={(ev) => {
            // Enter sends: this is a one-line instruction far more often than a paragraph, and shift-enter
            // is still there for the times it isn't.
            if (ev.key === 'Enter' && !ev.shiftKey) {
              ev.preventDefault();
              void ask();
            }
          }}
          rows={2}
          spellCheck={false}
          disabled={disabled}
          placeholder="e.g. give my laptop SSH access to everything tagged tag:server"
          className="w-full resize-y rounded-lg border border-white/10 bg-black/40 px-3 py-2 text-sm leading-relaxed text-zinc-100 outline-none transition-colors placeholder:text-zinc-600 focus:border-primary/50 disabled:opacity-50"
        />

        <div className="flex flex-wrap items-center gap-2">
          {!prompt.trim() &&
            EXAMPLES.map((example) => (
              <button
                key={example}
                type="button"
                onClick={() => setPrompt(example)}
                disabled={disabled}
                className="cursor-pointer rounded-full border border-white/10 px-2.5 py-1 text-[11px] text-zinc-500 transition-colors hover:border-white/20 hover:text-zinc-300 disabled:opacity-40"
              >
                {example}
              </button>
            ))}
          <div className="ml-auto">
            <Button variant="primary" onClick={() => void ask()} disabled={!prompt.trim() || assist.isPending}>
              {assist.isPending ? <Loader2 className="h-3.5 w-3.5 animate-spin" /> : <Wand2 className="h-3.5 w-3.5" />}
              {assist.isPending ? 'Drafting…' : 'Ask'}
            </Button>
          </div>
        </div>

        {assist.error && <ErrorNote>{assistFailure(assist.error)}</ErrorNote>}
      </div>

      {proposal && (
        <div className="border-t border-white/10">
          {proposal.explanation && (
            <p className="px-3 py-2.5 text-xs leading-relaxed whitespace-pre-wrap text-zinc-300">
              {proposal.explanation}
            </p>
          )}
          <div className="border-t border-white/10">
            <DiffBody before={policy} after={proposal.policy} />
          </div>
          <div className="flex items-center gap-2 border-t border-white/10 px-3 py-2">
            <span className="text-[11px] text-zinc-600">Applying only fills the editor — you still press Save.</span>
            <span className="ml-auto flex items-center gap-2">
              <Button onClick={() => setProposal(null)}>
                <X className="h-3.5 w-3.5" />
                Discard
              </Button>
              <Button variant="primary" onClick={apply}>
                <Check className="h-3.5 w-3.5" />
                Apply to editor
              </Button>
            </span>
          </div>
        </div>
      )}
    </Card>
  );
};