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

/**
 * useComplianceAnalysis — compliance analysis state wired to PageStateContext.
 *
 * State is stored in the global context so it persists when the user navigates
 * to another module and returns. The `run` and `reset` actions are identical
 * to the previous hook API so CompliancePage needs no structural changes.
 */

import { useCallback } from 'react';
import { usePageState } from '../../contexts';
import type {
  ComplianceMeta,
  ComplianceState,
  ComplianceSourceEvent,
  ComplianceFindingEvent,
  ComplianceDonePayload,
  ComplianceConflict,
} from '../../contexts';

export type { ComplianceMeta, ComplianceState, ComplianceSourceEvent as SourceEvent, ComplianceFindingEvent as FindingEvent, ComplianceDonePayload as DonePayload, ComplianceConflict };
export type { ComplianceActionItem as ActionItem } from '../../contexts';
export type AnalysisStatus = import('../../contexts').ComplianceStatus;
export type AnalysisMeta = ComplianceMeta;

const TOKEN_KEY = 'auth_token';
function authHeader(): Record<string, string> {
  const t = localStorage.getItem(TOKEN_KEY);
  return t ? { Authorization: `Bearer ${t}` } : {};
}

const INITIAL_STATE: ComplianceState = {
  status: 'idle',
  stageLabel: '',
  stageKey: '',
  meta: null,
  sources: [],
  findings: [],
  done: null,
  errorText: '',
  analysisId: null,
  isReadOnly: false,
  progress: null,
  conflicts: [],
};

export function useComplianceAnalysis() {
  const { complianceState: state, setComplianceState: setState, complianceAbortRef, resetCompliance: reset } = usePageState();

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

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

    try {
      const res = await fetch('/api/v1/compliance/analyze-stream', {
        method: 'POST',
        headers: authHeader(),
        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 === 'progress') {
              // Real per-clause progress update from backend
              setState(s => ({ ...s, progress: { done: j.done ?? 0, total: j.total ?? 0 } }));
            } else if (j.type === 'source') {
              const src: ComplianceSourceEvent = {
                standard: j.standard ?? '',
                clause: j.clause ?? '',
                score: j.score ?? 0,
                status: j.status ?? 'retrieved',
                full_content: j.full_content ?? '',
                clause_index: j.clause_index,
              };
              setState(s => ({ ...s, sources: [...s.sources, src] }));
            } else if (j.type === 'finding') {
              const finding: ComplianceFindingEvent = {
                title: j.title ?? '',
                desc: j.desc ?? '',
                status: j.status ?? 'info',
                clause_ref: j.clause_ref,
                confidence: j.confidence,
                source_refs: j.source_refs,
              };
              setState(s => ({ ...s, findings: [...s.findings, finding] }));
            } else if (j.type === 'conflicts') {
              // Cross-clause conflicts detected after all findings finish
              setState(s => ({ ...s, conflicts: j.items ?? [] }));
            } else if (j.type === 'done') {
              const payload: ComplianceDonePayload = {
                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 === 'saved') {
              setState(s => ({ ...s, analysisId: j.analysis_id ?? null }));
            } 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) }));
    }
  }, [setState, complianceAbortRef]);

  return { state, run, reset };
}