/**
 * LLM Store
 * Zustand store for LLM Agent session and message management
 * Handles session lifecycle, message history, and token tracking
 */

import { create } from 'zustand';
import { devtools, persist } from 'zustand/middleware';
import { apiClient } from '@/lib/apiClient';

// ============================================================================
// Types
// ============================================================================

export type MessageRole = 'user' | 'assistant' | 'system' | 'tool';

export interface LLMMessage {
  id: string;
  role: MessageRole;
  content: string;
  timestamp: string;
  modelName?: string;
  promptTokens?: number;
  completionTokens?: number;
  totalTokens?: number;
  toolCalls?: Array<{
    tool: string;
    params: Record<string, unknown>;
    result?: unknown;
  }>;
  metadata?: Record<string, unknown>;
}

export interface LLMSession {
  id: string;
  userId: string;
  title?: string;
  sessionType: 'general' | 'trading_advice' | 'education' | 'market_analysis' | 'support';
  status: 'active' | 'archived' | 'ended';
  messages: LLMMessage[];
  totalTokensUsed: number;
  relatedSymbols: string[];
  createdAt: string;
  updatedAt: string;
  endedAt?: string;
}

export interface TokenUsage {
  promptTokens: number;
  completionTokens: number;
  totalTokens: number;
  sessionTokens: number;
  dailyTokensUsed: number;
  dailyLimit: number;
}

// ============================================================================
// State Interface
// ============================================================================

interface LLMState {
  // Sessions
  sessions: LLMSession[];
  activeSession: LLMSession | null;
  activeSessionId: string | null;

  // Messages (for active session)
  messages: LLMMessage[];

  // Token tracking
  tokenUsage: TokenUsage;

  // Loading states
  loading: boolean;
  loadingSessions: boolean;
  loadingMessages: boolean;
  sendingMessage: boolean;

  // Error state
  error: string | null;

  // Service status
  llmServiceHealthy: boolean;

  // Actions - Session Lifecycle
  createSession: (sessionType?: LLMSession['sessionType'], title?: string) => Promise<LLMSession | null>;
  getSession: (sessionId: string) => Promise<LLMSession | null>;
  loadSessions: () => Promise<void>;
  setActiveSession: (sessionId: string | null) => void;
  endSession: (sessionId: string) => Promise<boolean>;
  archiveSession: (sessionId: string) => Promise<boolean>;
  deleteSession: (sessionId: string) => Promise<boolean>;

  // Actions - Messages
  sendMessage: (content: string, context?: Record<string, unknown>) => Promise<LLMMessage | null>;
  getHistory: (sessionId?: string, limit?: number, offset?: number) => Promise<LLMMessage[]>;
  clearMessages: () => void;

  // Actions - Token Management
  getTokenUsage: () => Promise<TokenUsage | null>;
  refreshTokenUsage: () => Promise<void>;

  // Actions - Utility
  checkHealth: () => Promise<boolean>;
  clearError: () => void;
  reset: () => void;
}

// ============================================================================
// Initial State
// ============================================================================

const initialTokenUsage: TokenUsage = {
  promptTokens: 0,
  completionTokens: 0,
  totalTokens: 0,
  sessionTokens: 0,
  dailyTokensUsed: 0,
  dailyLimit: 100000,
};

const initialState = {
  sessions: [],
  activeSession: null,
  activeSessionId: null,
  messages: [],
  tokenUsage: initialTokenUsage,
  loading: false,
  loadingSessions: false,
  loadingMessages: false,
  sendingMessage: false,
  error: null,
  llmServiceHealthy: false,
};

// ============================================================================
// Store
// ============================================================================

export const useLLMStore = create<LLMState>()(
  devtools(
    persist(
      (set, get) => ({
        ...initialState,

        // ======================================================================
        // Session Lifecycle Actions
        // ======================================================================

        createSession: async (
          sessionType: LLMSession['sessionType'] = 'general',
          title?: string
        ) => {
          set({ loading: true, error: null }, false, 'llm/createSession/start');

          try {
            const response = await apiClient.post('/proxy/llm/sessions', {
              sessionType,
              title,
            });

            const newSession: LLMSession = {
              id: response.data.id || response.data.sessionId,
              userId: response.data.userId || '',
              title: title || `Session ${new Date().toLocaleDateString()}`,
              sessionType,
              status: 'active',
              messages: [],
              totalTokensUsed: 0,
              relatedSymbols: [],
              createdAt: response.data.createdAt || new Date().toISOString(),
              updatedAt: response.data.createdAt || new Date().toISOString(),
            };

            set(
              (state) => ({
                sessions: [newSession, ...state.sessions],
                activeSession: newSession,
                activeSessionId: newSession.id,
                messages: [],
                loading: false,
              }),
              false,
              'llm/createSession/success'
            );

            localStorage.setItem('llmActiveSessionId', newSession.id);
            return newSession;
          } catch (error) {
            const message = error instanceof Error ? error.message : 'Failed to create session';
            set({ error: message, loading: false }, false, 'llm/createSession/error');
            console.error('Error creating LLM session:', error);
            return null;
          }
        },

        getSession: async (sessionId: string) => {
          set({ loading: true, error: null }, false, 'llm/getSession/start');

          try {
            const response = await apiClient.get(`/proxy/llm/sessions/${sessionId}`);
            const session = response.data as LLMSession;

            set(
              (state) => ({
                sessions: state.sessions.map((s) =>
                  s.id === sessionId ? session : s
                ),
                activeSession: session,
                activeSessionId: session.id,
                messages: session.messages || [],
                loading: false,
              }),
              false,
              'llm/getSession/success'
            );

            localStorage.setItem('llmActiveSessionId', sessionId);
            return session;
          } catch (error) {
            const message = error instanceof Error ? error.message : 'Failed to get session';
            set({ error: message, loading: false }, false, 'llm/getSession/error');
            console.error('Error getting LLM session:', error);
            return null;
          }
        },

        loadSessions: async () => {
          set({ loadingSessions: true, error: null }, false, 'llm/loadSessions/start');

          try {
            const response = await apiClient.get('/proxy/llm/sessions');
            const sessions = (response.data || []) as LLMSession[];

            set({ sessions, loadingSessions: false }, false, 'llm/loadSessions/success');

            // Restore last active session if available
            const savedSessionId = localStorage.getItem('llmActiveSessionId');
            if (savedSessionId) {
              const savedSession = sessions.find((s) => s.id === savedSessionId);
              if (savedSession && savedSession.status === 'active') {
                get().getSession(savedSessionId);
              }
            }
          } catch (error) {
            const message = error instanceof Error ? error.message : 'Failed to load sessions';
            set({ error: message, loadingSessions: false }, false, 'llm/loadSessions/error');
            console.error('Error loading LLM sessions:', error);
          }
        },

        setActiveSession: (sessionId: string | null) => {
          if (!sessionId) {
            set(
              {
                activeSession: null,
                activeSessionId: null,
                messages: [],
              },
              false,
              'llm/setActiveSession/clear'
            );
            localStorage.removeItem('llmActiveSessionId');
            return;
          }

          const session = get().sessions.find((s) => s.id === sessionId);
          if (session) {
            set(
              {
                activeSession: session,
                activeSessionId: sessionId,
                messages: session.messages || [],
              },
              false,
              'llm/setActiveSession'
            );
            localStorage.setItem('llmActiveSessionId', sessionId);
          } else {
            // Session not in local state, fetch it
            get().getSession(sessionId);
          }
        },

        endSession: async (sessionId: string) => {
          set({ loading: true, error: null }, false, 'llm/endSession/start');

          try {
            await apiClient.post(`/proxy/llm/sessions/${sessionId}/end`);

            set(
              (state) => {
                const updatedSessions = state.sessions.map((s) =>
                  s.id === sessionId
                    ? { ...s, status: 'ended' as const, endedAt: new Date().toISOString() }
                    : s
                );

                const isActive = state.activeSessionId === sessionId;
                return {
                  sessions: updatedSessions,
                  activeSession: isActive ? null : state.activeSession,
                  activeSessionId: isActive ? null : state.activeSessionId,
                  messages: isActive ? [] : state.messages,
                  loading: false,
                };
              },
              false,
              'llm/endSession/success'
            );

            if (get().activeSessionId === null) {
              localStorage.removeItem('llmActiveSessionId');
            }

            return true;
          } catch (error) {
            const message = error instanceof Error ? error.message : 'Failed to end session';
            set({ error: message, loading: false }, false, 'llm/endSession/error');
            console.error('Error ending LLM session:', error);
            return false;
          }
        },

        archiveSession: async (sessionId: string) => {
          set({ loading: true, error: null }, false, 'llm/archiveSession/start');

          try {
            await apiClient.post(`/proxy/llm/sessions/${sessionId}/archive`);

            set(
              (state) => ({
                sessions: state.sessions.map((s) =>
                  s.id === sessionId ? { ...s, status: 'archived' as const } : s
                ),
                activeSession:
                  state.activeSessionId === sessionId ? null : state.activeSession,
                activeSessionId:
                  state.activeSessionId === sessionId ? null : state.activeSessionId,
                messages: state.activeSessionId === sessionId ? [] : state.messages,
                loading: false,
              }),
              false,
              'llm/archiveSession/success'
            );

            return true;
          } catch (error) {
            const message = error instanceof Error ? error.message : 'Failed to archive session';
            set({ error: message, loading: false }, false, 'llm/archiveSession/error');
            console.error('Error archiving LLM session:', error);
            return false;
          }
        },

        deleteSession: async (sessionId: string) => {
          set({ loading: true, error: null }, false, 'llm/deleteSession/start');

          try {
            await apiClient.delete(`/proxy/llm/sessions/${sessionId}`);

            set(
              (state) => {
                const newSessions = state.sessions.filter((s) => s.id !== sessionId);
                const isActive = state.activeSessionId === sessionId;

                return {
                  sessions: newSessions,
                  activeSession: isActive ? null : state.activeSession,
                  activeSessionId: isActive ? null : state.activeSessionId,
                  messages: isActive ? [] : state.messages,
                  loading: false,
                };
              },
              false,
              'llm/deleteSession/success'
            );

            if (get().activeSessionId === null) {
              localStorage.removeItem('llmActiveSessionId');
            }

            return true;
          } catch (error) {
            const message = error instanceof Error ? error.message : 'Failed to delete session';
            set({ error: message, loading: false }, false, 'llm/deleteSession/error');
            console.error('Error deleting LLM session:', error);
            return false;
          }
        },

        // ======================================================================
        // Message Actions
        // ======================================================================

        sendMessage: async (content: string, context?: Record<string, unknown>) => {
          const { activeSessionId } = get();

          // Auto-create session if none exists
          if (!activeSessionId) {
            const newSession = await get().createSession();
            if (!newSession) {
              return null;
            }
          }

          const sessionId = get().activeSessionId;
          if (!sessionId) {
            set({ error: 'No active session' }, false, 'llm/sendMessage/noSession');
            return null;
          }

          set({ sendingMessage: true, error: null }, false, 'llm/sendMessage/start');

          // Optimistically add user message
          const userMessage: LLMMessage = {
            id: `temp-${Date.now()}`,
            role: 'user',
            content,
            timestamp: new Date().toISOString(),
          };

          set(
            (state) => ({
              messages: [...state.messages, userMessage],
            }),
            false,
            'llm/sendMessage/optimistic'
          );

          try {
            const response = await apiClient.post(`/proxy/llm/sessions/${sessionId}/messages`, {
              content,
              context,
            });

            const assistantMessage: LLMMessage = {
              id: response.data.id || `msg-${Date.now()}`,
              role: 'assistant',
              content: response.data.content || response.data.message,
              timestamp: response.data.timestamp || new Date().toISOString(),
              modelName: response.data.modelName,
              promptTokens: response.data.promptTokens,
              completionTokens: response.data.completionTokens,
              totalTokens: response.data.totalTokens,
              toolCalls: response.data.toolCalls,
              metadata: response.data.metadata,
            };

            // Update user message with server ID if provided
            const serverUserMessage: LLMMessage = {
              ...userMessage,
              id: response.data.userMessageId || userMessage.id,
            };

            set(
              (state) => ({
                messages: [
                  ...state.messages.filter((m) => m.id !== userMessage.id),
                  serverUserMessage,
                  assistantMessage,
                ],
                tokenUsage: {
                  ...state.tokenUsage,
                  promptTokens: state.tokenUsage.promptTokens + (assistantMessage.promptTokens || 0),
                  completionTokens: state.tokenUsage.completionTokens + (assistantMessage.completionTokens || 0),
                  totalTokens: state.tokenUsage.totalTokens + (assistantMessage.totalTokens || 0),
                  sessionTokens: state.tokenUsage.sessionTokens + (assistantMessage.totalTokens || 0),
                },
                sendingMessage: false,
              }),
              false,
              'llm/sendMessage/success'
            );

            return assistantMessage;
          } catch (error) {
            // Remove optimistic message on error
            set(
              (state) => ({
                messages: state.messages.filter((m) => m.id !== userMessage.id),
                error: error instanceof Error ? error.message : 'Failed to send message',
                sendingMessage: false,
              }),
              false,
              'llm/sendMessage/error'
            );
            console.error('Error sending LLM message:', error);
            return null;
          }
        },

        getHistory: async (sessionId?: string, limit: number = 50, offset: number = 0) => {
          const targetSessionId = sessionId || get().activeSessionId;

          if (!targetSessionId) {
            return [];
          }

          set({ loadingMessages: true, error: null }, false, 'llm/getHistory/start');

          try {
            const response = await apiClient.get(
              `/proxy/llm/sessions/${targetSessionId}/messages`,
              {
                params: { limit, offset },
              }
            );

            const messages = (response.data || []) as LLMMessage[];

            if (!sessionId || sessionId === get().activeSessionId) {
              set({ messages, loadingMessages: false }, false, 'llm/getHistory/success');
            } else {
              set({ loadingMessages: false }, false, 'llm/getHistory/success');
            }

            return messages;
          } catch (error) {
            const message = error instanceof Error ? error.message : 'Failed to get message history';
            set({ error: message, loadingMessages: false }, false, 'llm/getHistory/error');
            console.error('Error getting LLM message history:', error);
            return [];
          }
        },

        clearMessages: () => {
          set({ messages: [] }, false, 'llm/clearMessages');
        },

        // ======================================================================
        // Token Management Actions
        // ======================================================================

        getTokenUsage: async () => {
          try {
            const response = await apiClient.get('/proxy/llm/tokens/usage');
            const usage = response.data as TokenUsage;

            set({ tokenUsage: usage }, false, 'llm/getTokenUsage/success');
            return usage;
          } catch (error) {
            console.error('Error getting token usage:', error);
            return null;
          }
        },

        refreshTokenUsage: async () => {
          await get().getTokenUsage();
        },

        // ======================================================================
        // Utility Actions
        // ======================================================================

        checkHealth: async () => {
          try {
            const response = await apiClient.get('/proxy/llm/health');
            const healthy = response.data?.status === 'healthy' || response.status === 200;

            set({ llmServiceHealthy: healthy }, false, 'llm/checkHealth');
            return healthy;
          } catch (error) {
            set({ llmServiceHealthy: false }, false, 'llm/checkHealth/error');
            console.error('Error checking LLM health:', error);
            return false;
          }
        },

        clearError: () => {
          set({ error: null }, false, 'llm/clearError');
        },

        reset: () => {
          set(initialState, false, 'llm/reset');
          localStorage.removeItem('llmActiveSessionId');
        },
      }),
      {
        name: 'llm-storage',
        partialize: (state) => ({
          activeSessionId: state.activeSessionId,
          tokenUsage: state.tokenUsage,
        }),
      }
    ),
    { name: 'LLMStore' }
  )
);

// ============================================================================
// Selectors (for performance optimization)
// ============================================================================

export const useLLMSessions = () => useLLMStore((state) => state.sessions);
export const useActiveSession = () => useLLMStore((state) => state.activeSession);
export const useActiveSessionId = () => useLLMStore((state) => state.activeSessionId);
export const useLLMMessages = () => useLLMStore((state) => state.messages);
export const useTokenUsage = () => useLLMStore((state) => state.tokenUsage);
export const useLLMLoading = () => useLLMStore((state) => state.loading);
export const useSendingMessage = () => useLLMStore((state) => state.sendingMessage);
export const useLLMError = () => useLLMStore((state) => state.error);
export const useLLMServiceHealthy = () => useLLMStore((state) => state.llmServiceHealthy);

export default useLLMStore;
