AIRegulation-DocAnalysis/frontend/src/pages/Compliance/useComplianceAnalysis.ts

import { useState, useCallback, useRef } from 'react';

export type AnalysisStatus = 'idle' | 'streaming' | 'done' | 'error';

export interface SourceEvent {
  standard: string;
  clause: string;
  score: number;
  status: string;
  full_content: string;
}

export interface FindingEvent {
  title: string;
  desc: string;
  status: 'ok' | 'warn' | 'risk';
  clause_ref?: string;
}

export interface ActionItem {
  label: string;
  value: string;
  risk?: boolean;
}

export interface DonePayload {
  conclusion: string;
  actions: ActionItem[];
  risk_score: number;
  highlight_terms: string[];
  para_text: string;
}

export interface AnalysisMeta {
  title: string;
  sourceType: 'text' | 'doc' | 'upload';
  startedAt: string; // ISO timestamp
}

export interface AnalysisState {
  status: AnalysisStatus;
  stageLabel: string;
  stageKey: string;
  meta: AnalysisMeta | null;
  sources: SourceEvent[];
  findings: FindingEvent[];
  done: DonePayload | null;
  errorText: string;
}

const INITIAL_STATE: AnalysisState = {
  status: 'idle',
  stageLabel: '',
  stageKey: '',
  meta: null,
  sources: [],
  findings: [],
  done: null,
  errorText: '',
};

export function useComplianceAnalysis() {
  const [state, setState] = useState<AnalysisState>(INITIAL_STATE);
  const abortRef = useRef<AbortController | null>(null);

  const reset = useCallback(() => {
    abortRef.current?.abort();
    setState(INITIAL_STATE);
  }, []);

  const run = useCallback(async (formData: FormData, meta: AnalysisMeta) => {
    abortRef.current?.abort();
    const ctrl = new AbortController();
    abortRef.current = ctrl;

    setState({ ...INITIAL_STATE, status: 'streaming', stageLabel: 'Starting…', meta });

    try {
      const res = await fetch('/api/v1/compliance/analyze-stream', {
        method: 'POST',
        body: formData,
        signal: ctrl.signal,
      });

      if (!res.ok) {
        const txt = await res.text();
        setState(s => ({ ...s, status: 'error', errorText: `HTTP ${res.status}: ${txt}` }));
        return;
      }

      if (!res.body) {
        setState(s => ({ ...s, status: 'error', errorText: 'No response stream' }));
        return;
      }

      const reader = res.body.getReader();
      const dec = new TextDecoder();
      let buffer = '';

      while (true) {
        const { done, value } = await reader.read();
        if (done) break;
        buffer += dec.decode(value, { stream: true });

        const blocks = buffer.split('\n\n');
        buffer = blocks.pop() ?? '';

        for (const block of blocks) {
          const dataLine = block.split('\n').find(l => l.startsWith('data: '));
          if (!dataLine) continue;
          const raw = dataLine.slice(6).trim();
          if (!raw) continue;

          try {
            const j = JSON.parse(raw);

            if (j.type === 'stage') {
              setState(s => ({ ...s, stageLabel: j.label ?? '', stageKey: j.stage ?? '' }));
            } else if (j.type === 'source') {
              const src: SourceEvent = {
                standard: j.standard ?? '',
                clause: j.clause ?? '',
                score: j.score ?? 0,
                status: j.status ?? 'retrieved',
                full_content: j.full_content ?? '',
              };
              setState(s => ({ ...s, sources: [...s.sources, src] }));
            } else if (j.type === 'finding') {
              const finding: FindingEvent = {
                title: j.title ?? '',
                desc: j.desc ?? '',
                status: j.status ?? 'info',
                clause_ref: j.clause_ref,
              };
              setState(s => ({ ...s, findings: [...s.findings, finding] }));
            } else if (j.type === 'done') {
              const payload: DonePayload = {
                conclusion: j.conclusion ?? '',
                actions: j.actions ?? [],
                risk_score: j.risk_score ?? 0,
                highlight_terms: j.highlight_terms ?? [],
                para_text: j.para_text ?? '',
              };
              setState(s => ({ ...s, status: 'done', done: payload, stageKey: 'concluding', stageLabel: 'Complete' }));
            } else if (j.type === 'error') {
              setState(s => ({ ...s, status: 'error', errorText: j.text ?? 'Unknown error' }));
            }
          } catch { /* skip malformed */ }
        }
      }

      // Mark done if stream ended without explicit done event
      setState(s => s.status === 'streaming' ? { ...s, status: 'done', stageKey: 'concluding', stageLabel: 'Complete' } : s);
    } catch (e: unknown) {
      if (e instanceof Error && e.name === 'AbortError') return;
      setState(s => ({ ...s, status: 'error', errorText: String(e) }));
    }
  }, []);

  return { state, run, reset };
}