import { create } from 'zustand';
import { persist } from 'zustand/middleware';
import { getGatewayClient, AgentStreamDelta } from '../lib/gateway-client';

export interface Message {
  id: string;
  role: 'user' | 'assistant' | 'tool';
  content: string;
  timestamp: Date;
  streaming?: boolean;
  toolName?: string;
  toolInput?: string;
  toolOutput?: string;
  error?: string;
}

export interface Conversation {
  id: string;
  title: string;
  messages: Message[];
  sessionKey: string | null;
  createdAt: Date;
  updatedAt: Date;
}

export interface Agent {
  id: string;
  name: string;
  icon: string;
  color: string;
  lastMessage: string;
  time: string;
}

interface ChatState {
  messages: Message[];
  conversations: Conversation[];
  currentConversationId: string | null;
  agents: Agent[];
  currentAgent: Agent | null;
  isStreaming: boolean;
  currentModel: string;
  sessionKey: string | null;

  addMessage: (message: Message) => void;
  updateMessage: (id: string, updates: Partial<Message>) => void;
  setCurrentAgent: (agent: Agent) => void;
  setCurrentModel: (model: string) => void;
  sendMessage: (content: string) => Promise<void>;
  initStreamListener: () => () => void;
  newConversation: () => void;
  switchConversation: (id: string) => void;
  deleteConversation: (id: string) => void;
}

function generateConvId(): string {
  return `conv_${Date.now()}_${Math.random().toString(36).slice(2, 6)}`;
}

function deriveTitle(messages: Message[]): string {
  const firstUser = messages.find(m => m.role === 'user');
  if (firstUser) {
    const text = firstUser.content.trim();
    return text.length > 30 ? text.slice(0, 30) + '...' : text;
  }
  return '新对话';
}

export const useChatStore = create<ChatState>()(
  persist(
    (set, get) => ({
  messages: [],
  conversations: [],
  currentConversationId: null,
  agents: [
    {
      id: '1',
      name: 'ZCLAW',
      icon: '🦞',
      color: 'bg-gradient-to-br from-orange-500 to-red-500',
      lastMessage: '发送消息开始对话',
      time: '',
    },
  ],
  currentAgent: null,
  isStreaming: false,
  currentModel: 'glm-5',
  sessionKey: null,

  addMessage: (message) =>
    set((state) => ({ messages: [...state.messages, message] })),

  updateMessage: (id, updates) =>
    set((state) => ({
      messages: state.messages.map((m) =>
        m.id === id ? { ...m, ...updates } : m
      ),
    })),

  setCurrentAgent: (agent) => set({ currentAgent: agent }),

  setCurrentModel: (model) => set({ currentModel: model }),

  newConversation: () => {
    const state = get();
    let conversations = [...state.conversations];

    // Save current conversation if it has messages
    if (state.messages.length > 0) {
      const currentId = state.currentConversationId || generateConvId();
      const existingIdx = conversations.findIndex(c => c.id === currentId);
      const conv: Conversation = {
        id: currentId,
        title: deriveTitle(state.messages),
        messages: [...state.messages],
        sessionKey: state.sessionKey,
        createdAt: existingIdx >= 0 ? conversations[existingIdx].createdAt : new Date(),
        updatedAt: new Date(),
      };
      if (existingIdx >= 0) {
        conversations[existingIdx] = conv;
      } else {
        conversations = [conv, ...conversations];
      }
    }

    set({
      conversations,
      messages: [],
      sessionKey: null,
      isStreaming: false,
      currentConversationId: null,
    });
  },

  switchConversation: (id: string) => {
    const state = get();
    let conversations = [...state.conversations];

    // Save current conversation first
    if (state.messages.length > 0 && state.currentConversationId) {
      const existingIdx = conversations.findIndex(c => c.id === state.currentConversationId);
      if (existingIdx >= 0) {
        conversations[existingIdx] = {
          ...conversations[existingIdx],
          messages: [...state.messages],
          sessionKey: state.sessionKey,
          updatedAt: new Date(),
          title: deriveTitle(state.messages),
        };
      }
    }

    const target = conversations.find(c => c.id === id);
    if (target) {
      set({
        conversations,
        messages: [...target.messages],
        sessionKey: target.sessionKey,
        currentConversationId: target.id,
        isStreaming: false,
      });
    }
  },

  deleteConversation: (id: string) => {
    const state = get();
    const conversations = state.conversations.filter(c => c.id !== id);
    if (state.currentConversationId === id) {
      set({ conversations, messages: [], sessionKey: null, currentConversationId: null, isStreaming: false });
    } else {
      set({ conversations });
    }
  },

  sendMessage: async (content: string) => {
    const { addMessage, currentModel, sessionKey } = get();

    // Add user message
    const userMsg: Message = {
      id: `user_${Date.now()}`,
      role: 'user',
      content,
      timestamp: new Date(),
    };
    addMessage(userMsg);

    // Create placeholder assistant message for streaming
    const assistantId = `assistant_${Date.now()}`;
    const assistantMsg: Message = {
      id: assistantId,
      role: 'assistant',
      content: '',
      timestamp: new Date(),
      streaming: true,
    };
    addMessage(assistantMsg);
    set({ isStreaming: true });

    try {
      const client = getGatewayClient();
      const result = await client.chat(content, {
        sessionKey: sessionKey || undefined,
        model: currentModel,
      });

      // Store session key for continuity
      if (!sessionKey) {
        set({ sessionKey: `session_${Date.now()}` });
      }

      // The actual streaming content comes via the 'agent' event listener
      // set in initStreamListener(). The runId links events to this message.
      // Store runId on the message for correlation
      set((state) => ({
        messages: state.messages.map((m) =>
          m.id === assistantId ? { ...m, toolInput: result.runId } : m
        ),
      }));
    } catch (err: any) {
      // Gateway not connected — show error in the assistant bubble
      set((state) => ({
        isStreaming: false,
        messages: state.messages.map((m) =>
          m.id === assistantId
            ? {
                ...m,
                content: `⚠️ ${err.message || '无法连接 Gateway'}`,
                streaming: false,
                error: err.message,
              }
            : m
        ),
      }));
    }
  },

  initStreamListener: () => {
    const client = getGatewayClient();

    const unsubscribe = client.onAgentStream((delta: AgentStreamDelta) => {
      const state = get();

      // Find the currently streaming assistant message
      const streamingMsg = [...state.messages]
        .reverse()
        .find((m) => m.role === 'assistant' && m.streaming);

      if (!streamingMsg) return;

      if (delta.stream === 'assistant' && delta.delta) {
        // Append text delta to the streaming message
        set((s) => ({
          messages: s.messages.map((m) =>
            m.id === streamingMsg.id
              ? { ...m, content: m.content + delta.delta }
              : m
          ),
        }));
      } else if (delta.stream === 'tool') {
        // Add a tool message
        const toolMsg: Message = {
          id: `tool_${Date.now()}_${Math.random().toString(36).slice(2, 6)}`,
          role: 'tool',
          content: delta.toolOutput || '',
          timestamp: new Date(),
          toolName: delta.tool,
          toolInput: delta.toolInput,
          toolOutput: delta.toolOutput,
        };
        set((s) => ({ messages: [...s.messages, toolMsg] }));
      } else if (delta.stream === 'lifecycle') {
        if (delta.phase === 'end' || delta.phase === 'error') {
          // Mark streaming complete
          set((s) => ({
            isStreaming: false,
            messages: s.messages.map((m) =>
              m.id === streamingMsg.id
                ? {
                    ...m,
                    streaming: false,
                    error: delta.phase === 'error' ? delta.error : undefined,
                  }
                : m
            ),
          }));
        }
      }
    });

    return unsubscribe;
  },
}),
    {
      name: 'zclaw-chat-storage',
      partialize: (state) => ({
        conversations: state.conversations,
        currentModel: state.currentModel,
      }),
      onRehydrateStorage: () => (state) => {
        // Rehydrate Date objects from JSON strings
        if (state?.conversations) {
          for (const conv of state.conversations) {
            conv.createdAt = new Date(conv.createdAt);
            conv.updatedAt = new Date(conv.updatedAt);
            for (const msg of conv.messages) {
              msg.timestamp = new Date(msg.timestamp);
              msg.streaming = false; // Never restore streaming state
            }
          }
        }
      },
    },
  ),
);
