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

import { callCoreRpc } from '../services/coreRpcClient';
import { aiListMemoryFiles, type GraphRelation, memoryGraphQuery } from '../utils/tauriCommands';

export type AIStatus = 'idle' | 'initializing' | 'ready' | 'error';

const POLL_MS = 5000;

interface SessionEntry {
  sessionId: string;
  updatedAt: number;
  inputTokens: number;
  outputTokens: number;
  totalTokens: number;
  compactionCount: number;
  memoryFlushAt?: number;
}

interface SessionStats {
  total: number;
  totalTokens: number;
  compactions: number;
  memoryFlushes: number;
}

export interface IntelligenceStats {
  sessions: SessionStats | null;
  memoryFiles: number | null;
  entities: Record<string, number> | null;
  entityError: boolean;
  aiStatus: AIStatus;
  isLoading: boolean;
  refetch: () => void;
}

/** Derive entity-type counts from local graph relations. */
function entityCountsFromRelations(relations: GraphRelation[]): Record<string, number> {
  const counts: Record<string, number> = {};
  for (const rel of relations) {
    const types = (rel.attrs?.entity_types ?? {}) as Record<string, string>;
    const subjectType = types.subject ?? 'entity';
    const objectType = types.object ?? 'entity';
    counts[subjectType] = (counts[subjectType] ?? 0) + 1;
    counts[objectType] = (counts[objectType] ?? 0) + 1;
  }
  return counts;
}

export function useIntelligenceStats(): IntelligenceStats {
  const [aiStatus, setAiStatus] = useState<AIStatus>('idle');
  const [sessions, setSessions] = useState<SessionStats | null>(null);
  const [memoryFiles, setMemoryFiles] = useState<number | null>(null);
  const [entities, setEntities] = useState<Record<string, number> | null>(null);
  const [entityError, setEntityError] = useState(false);
  const [isLoading, setIsLoading] = useState(true);

  const fetchStats = useCallback(async () => {
    setAiStatus('initializing');
    setIsLoading(true);
    let hasSuccess = false;

    // Fetch local stats (Tauri invoke)
    try {
      const index = await callCoreRpc<Record<string, SessionEntry>>({
        method: 'ai.sessions_load_index',
      });
      const entries = Object.values(index);
      setSessions({
        total: entries.length,
        totalTokens: entries.reduce((sum, e) => sum + (e.totalTokens || 0), 0),
        compactions: entries.reduce((sum, e) => sum + (e.compactionCount || 0), 0),
        memoryFlushes: entries.filter(e => e.memoryFlushAt).length,
      });
      hasSuccess = true;
    } catch {
      setSessions(null);
    }

    try {
      // Empty string lists the memory root; the resolver joins it
      // onto `<workspace>/memory/`, so passing 'memory' here would
      // double up to `<workspace>/memory/memory` and miss the dir.
      const files = await aiListMemoryFiles('');
      setMemoryFiles(files.length);
      hasSuccess = true;
    } catch {
      setMemoryFiles(null);
    }

    // Derive entity counts from local graph store
    try {
      const relations = await memoryGraphQuery();
      const counts = entityCountsFromRelations(relations);
      if (Object.keys(counts).length > 0) {
        setEntities(counts);
        setEntityError(false);
      } else {
        setEntities(null);
        setEntityError(false);
      }
      hasSuccess = true;
    } catch {
      setEntities(null);
      setEntityError(true);
    }

    setAiStatus(hasSuccess ? 'ready' : 'error');
    setIsLoading(false);
  }, []);

  useEffect(() => {
    void fetchStats();
    const intervalId = window.setInterval(() => {
      void fetchStats();
    }, POLL_MS);

    return () => {
      window.clearInterval(intervalId);
    };
  }, [fetchStats]);

  return { sessions, memoryFiles, entities, entityError, aiStatus, isLoading, refetch: fetchStats };
}
