src/client/useAgent.ts

import { useCallback, useEffect, useMemo, useReducer, useRef, useState } from 'react';
import type { WorkerEvent } from '../shared/protocol';
import type { AgentMessage, AssistantMessage, DeviceInfo, ShellResult } from '../shared/types';
import { MODEL_TOTAL_BYTES } from '../worker/manifest';
import { WorkerError, getClient } from './agentClient';

// ───────────────────────── Trạng thái ─────────────────────────

export interface LoadState {
  phase: 'download' | 'compile' | 'warmup';
  loaded: number;
  total: number;
  cached: number;
  file?: string;
}

export type MeterTone = 'idle' | 'loading' | 'active' | 'error';

export interface Meter {
  /** Nơi việc đang chạy: GPU cho suy luận, JavaScript cho tool/shell. */
  engine: 'WebGPU' | 'JavaScript';
  tone: MeterTone;
  label: string;
  rate: number | null;
  samples: number[];
}

export interface ShellEntry {
  id: number;
  command: string;
  result: ShellResult | { stdout: ''; stderr: string; exitCode: 1 };
}

export interface AgentState {
  initialized: boolean;
  messages: AgentMessage[];
  streaming: AssistantMessage | null;
  files: Record<string, string>;
  busy: string | null;
  modelLoaded: boolean;
  load: LoadState | null;
  meter: Meter;
  notice: string;
  error: string;
  device: DeviceInfo | null;
  context: { input: number; output: number };
}

const IDLE_METER: Meter = { engine: 'WebGPU', tone: 'idle', label: 'Idle', rate: null, samples: [] };
const MAX_SAMPLES = 24;

const initialState: AgentState = {
  initialized: false,
  messages: [],
  streaming: null,
  files: {},
  busy: null,
  modelLoaded: false,
  load: null,
  meter: IDLE_METER,
  notice: '',
  error: '',
  device: null,
  context: { input: 0, output: 0 },
};

type Action = { type: 'event'; event: WorkerEvent } | { type: 'set'; patch: Partial<AgentState> };

const percent = (loaded: number, total: number) => Math.min(100, Math.floor((loaded / Math.max(total, 1)) * 100));

function reduce(state: AgentState, action: Action): AgentState {
  if (action.type === 'set') return { ...state, ...action.patch };
  const e = action.event;
  switch (e.type) {
    case 'initialized':
      return { ...state, initialized: true, messages: e.messages, files: e.files };
    case 'messages':
      return { ...state, messages: e.messages };
    case 'stream':
      return { ...state, streaming: e.message };
    case 'workspace':
      return { ...state, files: e.files };
    case 'busy': {
      const meter: Meter =
        e.action === 'load'
          ? { ...IDLE_METER, tone: 'loading', label: e.cachedOnly ? 'Loading local model…' : 'Loading…' }
          : e.action === 'prompt'
            ? { ...IDLE_METER, tone: 'active', label: 'Reading prompt…' }
            : e.action === 'shell'
              ? { ...IDLE_METER, engine: 'JavaScript', tone: 'active', label: 'Running bash…' }
              : { ...IDLE_METER, tone: 'active', label: 'Working…' };
      return { ...state, busy: e.action, meter, error: '' };
    }
    case 'idle':
      return {
        ...state,
        busy: null,
        load: null,
        streaming: null,
        meter: state.meter.tone === 'error' ? state.meter : { ...IDLE_METER, label: state.modelLoaded ? 'Ready' : 'Idle' },
      };
    case 'load_progress': {
      const load: LoadState = {
        phase: e.phase,
        loaded: e.loaded ?? 0,
        total: e.total ?? MODEL_TOTAL_BYTES,
        cached: e.cached ?? 0,
        file: e.file,
      };
      const label =
        e.phase === 'download'
          ? load.loaded < load.total
            ? `Downloading ${percent(load.loaded, load.total)}%`
            : 'Verifying…'
          : e.phase === 'compile'
            ? 'Preparing GPU…'
            : 'Warming up…';
      return { ...state, load, meter: { ...state.meter, tone: 'loading', label } };
    }
    case 'device':
      return { ...state, device: e.device };
    case 'loaded':
      return { ...state, modelLoaded: true, device: e.device, load: null, meter: { ...IDLE_METER, label: 'Ready' } };
    case 'fatal':
      return { ...state, modelLoaded: false, error: e.error, meter: { ...IDLE_METER, tone: 'error', label: 'Unavailable' } };
    case 'notice':
      return { ...state, notice: e.text };
    case 'context_usage':
      return { ...state, context: { input: e.inputTokens, output: e.outputTokens } };
    case 'context_trim':
      return {
        ...state,
        notice: `Dropped ${e.dropped} earlier message${e.dropped === 1 ? '' : 's'} to fit the 8,192-token context.`,
      };
    case 'inference_activity': {
      if (e.phase === 'prefill') return { ...state, meter: { ...IDLE_METER, tone: 'active', label: 'Reading prompt…' } };
      if (e.phase === 'end') return { ...state, meter: { ...IDLE_METER, tone: 'active', label: 'Working…' } };
      const samples = e.rate ? [...state.meter.samples, e.rate].slice(-MAX_SAMPLES) : state.meter.samples;
      return {
        ...state,
        meter: {
          engine: 'WebGPU',
          tone: 'active',
          label: e.rate ? `${e.rate.toFixed(1)} tok/s` : 'Generating…',
          rate: e.rate ?? null,
          samples,
        },
      };
    }
    case 'tool_start':
      return { ...state, meter: { ...IDLE_METER, engine: 'JavaScript', tone: 'active', label: `Running ${e.name}…` } };
    case 'tool_end':
      return { ...state, meter: { ...IDLE_METER, tone: 'active', label: 'Working…' } };
    case 'persistence_error':
      return { ...state, notice: e.error };
    default:
      return state;
  }
}

// ───────────────────────── Hook ─────────────────────────

export interface PendingSend {
  text: string;
  error?: string;
  onAccepted?: () => void;
}

export function useAgent() {
  const client = useMemo(() => getClient(), []);
  const [state, dispatch] = useReducer(reduce, initialState);
  const [pending, setPending] = useState<PendingSend | null>(null);
  const [shell, setShell] = useState<ShellEntry[]>([]);
  const loadedRef = useRef(false);
  const shellId = useRef(1);
  loadedRef.current = state.modelLoaded;

  useEffect(() => {
    const off = client.subscribe((event) => dispatch({ type: 'event', event }));
    client.ready.then((init) => dispatch({ type: 'event', event: init }));
    return off;
  }, [client]);

  const set = useCallback((patch: Partial<AgentState>) => dispatch({ type: 'set', patch }), []);

  const loadModel = useCallback(
    async (cachedOnly = false) => {
      set({ error: '' });
      await client.ready;
      await client.call('load', { cachedOnly });
    },
    [client, set],
  );

  /** Bắt đầu một lượt hội thoại. Trả về true nếu đã gửi ngay, false nếu cần người dùng xác nhận tải model. */
  const send = useCallback(
    async (text: string, onAccepted?: () => void): Promise<boolean> => {
      const trimmed = text.trim();
      if (!trimmed) return false;
      set({ error: '', notice: '' });
      await client.ready;

      if (!loadedRef.current) {
        let cacheError: string | undefined;
        if (await client.call<boolean>('cache_status')) {
          try {
            await loadModel(true);
          } catch (err) {
            if ((err as Error).name === 'AbortError') return false;
            if ((err as WorkerError).code !== 'MODEL_CACHE_MISS') cacheError = (err as Error).message;
          }
        }
        if (!loadedRef.current) {
          setPending({ text: trimmed, error: cacheError, onAccepted });
          return false;
        }
      }

      try {
        onAccepted?.();
        await client.call('prompt', { text: trimmed });
      } catch (err) {
        if ((err as Error).message !== 'Stopped.') set({ error: (err as Error).message });
      }
      return true;
    },
    [client, loadModel, set, setPending],
  );

  const confirmPending = useCallback(async () => {
    if (!pending) return;
    const { text, onAccepted } = pending;
    setPending({ text });
    try {
      await loadModel(false);
    } catch (err) {
      if ((err as Error).name === 'AbortError') {
        setPending(null);
        return;
      }
      setPending({ text, error: (err as Error).message, onAccepted });
      return;
    }
    setPending(null);
    try {
      onAccepted?.();
      await client.call('prompt', { text });
    } catch (err) {
      if ((err as Error).message !== 'Stopped.') set({ error: (err as Error).message });
    }
  }, [client, loadModel, pending, set, setPending]);

  const cancelPending = useCallback(() => {
    // Nếu đang tải, dừng luôn; bản nháp vẫn giữ nguyên trong ô soạn thảo.
    void client.call('stop').catch(() => {});
    setPending(null);
  }, [client, setPending]);

  const stop = useCallback(() => void client.call('stop').catch(() => {}), [client]);

  const newChat = useCallback(async () => {
    try {
      await client.call('new_chat');
      set({ notice: '', error: '', context: { input: 0, output: 0 } });
    } catch (err) {
      set({ error: (err as Error).message });
    }
  }, [client, set]);

  const writeFile = useCallback(
    async (path: string, content: string) => {
      try {
        await client.call('write', { path, content });
        return true;
      } catch (err) {
        set({ error: (err as Error).message });
        return false;
      }
    },
    [client, set],
  );

  const runShell = useCallback(
    async (command: string) => {
      const id = shellId.current++;
      try {
        const result = await client.call<ShellResult>('shell', { command });
        setShell((list) => [...list, { id, command, result }]);
      } catch (err) {
        setShell((list) => [
          ...list,
          { id, command, result: { stdout: '', stderr: (err as Error).message + '\n', exitCode: 1 } },
        ]);
      }
    },
    [client, setShell],
  );

  const clearShell = useCallback(() => setShell([]), [setShell]);
  const dismissNotice = useCallback(() => set({ notice: '', error: '' }), [set]);
  const setError = useCallback((error: string) => set({ error }), [set]);

  return {
    state,
    pending,
    shell,
    actions: { loadModel, send, confirmPending, cancelPending, stop, newChat, writeFile, runShell, clearShell, dismissNotice, setError },
  };
}