import { create } from 'zustand'
import type { Room, Message, AgentInfo, SkillMeta, WsEvent, TodoItem } from './types'

interface TokenUsage {
  promptTokens: number
  completionTokens: number
  totalTokens: number
}

interface AppState {
  theme: 'light' | 'dark'
  user: {
    name: string
    avatar_color: string
  } | null
  onboardingCompleted: boolean
  rooms: Room[]
  activeRoomId: string | null
  messages: Record<string, Message[]>
  tasks: Record<string, string>
  workspace: Record<string, string[]>
  agents: AgentInfo[]
  skills: SkillMeta[]
  page: 'chat' | 'agents' | 'skills' | 'market' | 'settings'
  ws: Record<string, WebSocket>
  tokenUsage: Record<string, TokenUsage>
  todoItems: Record<string, TodoItem[]>
  fileReaders: Record<string, Record<string, string[]>>
  workingFiles: Record<string, { agent: string; filename: string; title?: string }[]>
  workStartedAt: Record<string, number>  // roomId -> timestamp (ms)

  setTheme: (theme: 'light' | 'dark') => void
  toggleTheme: () => void
  setOnboardingCompleted: (completed: boolean) => void
  fetchUser: () => Promise<void>
  setPage: (p: AppState['page']) => void
  setActiveRoom: (id: string) => void
  fetchRooms: () => Promise<void>
  fetchAgents: () => Promise<void>
  fetchSkills: () => Promise<void>
  sendMessage: (roomId: string, content: string) => void
  connectRoom: (roomId: string) => void
  markFileRead: (roomId: string, filename: string) => void
  setRoomTeam: (roomId: string, team: string) => Promise<void>
}

const API = '/api'

export const useStore = create<AppState>((set, get) => {
  const getInitialTheme = (): 'light' | 'dark' => {
    if (typeof window !== 'undefined') {
      const saved = localStorage.getItem('theme')
      if (saved === 'light' || saved === 'dark') return saved
      return window.matchMedia('(prefers-color-scheme: dark)').matches ? 'dark' : 'light'
    }
    return 'dark'
  }

  const applyTheme = (theme: 'light' | 'dark') => {
    if (typeof window !== 'undefined') {
      document.documentElement.classList.remove('light', 'dark')
      document.documentElement.classList.add(theme)
      localStorage.setItem('theme', theme)
    }
  }

  const initialTheme = getInitialTheme()
  applyTheme(initialTheme)

  return {
    theme: initialTheme,
    user: null,
    onboardingCompleted: typeof window !== 'undefined' ? localStorage.getItem('onboarding_completed') === 'true' : false,
    rooms: [],
    activeRoomId: typeof window !== 'undefined' ? localStorage.getItem('activeRoomId') : null,
    messages: {},
    tasks: {},
    workspace: {},
    agents: [],
    skills: [],
    page: (typeof window !== 'undefined' ? localStorage.getItem('page') as AppState['page'] : null) || 'chat',
    ws: {},
    tokenUsage: {},
    todoItems: {},
    fileReaders: {},
    workingFiles: {},
    workStartedAt: {},

    setTheme: (theme) => {
      applyTheme(theme)
      set({ theme })
    },
    toggleTheme: () => {
      const newTheme = get().theme === 'dark' ? 'light' : 'dark'
      applyTheme(newTheme)
      set({ theme: newTheme })
    },
    setOnboardingCompleted: (completed) => {
      localStorage.setItem('onboarding_completed', completed ? 'true' : 'false')
      set({ onboardingCompleted: completed })
    },
    fetchUser: async () => {
      try {
        const res = await fetch(`${API}/user`)
        const data = await res.json()
        set({
          user: {
            name: data.name || '用户',
            avatar_color: data.avatar_color || '#5865F2'
          }
        })
      } catch {
        set({
          user: {
            name: '用户',
            avatar_color: '#5865F2'
          }
        })
      }
    },
    setPage: (page) => {
      localStorage.setItem('page', page)
      set({ page })
    },
    setActiveRoom: (id) => {
      localStorage.setItem('activeRoomId', id)
      set({ activeRoomId: id })
      get().connectRoom(id)
      fetch(`${API}/rooms/${id}/tasks`).then(r => r.json()).then(d => {
        set(s => ({ tasks: { ...s.tasks, [id]: d.content } }))
      })
      fetch(`${API}/rooms/${id}/workspace`).then(r => r.json()).then(files => {
        set(s => ({ workspace: { ...s.workspace, [id]: files || [] } }))
      })
      fetch(`${API}/rooms/${id}/messages`).then(r => r.json()).then(msgs => {
        if (msgs?.length) set(s => ({ messages: { ...s.messages, [id]: msgs } }))
      }).catch(() => {})
      fetch(`${API}/rooms/${id}/token-usage`).then(r => r.json()).then(data => {
        if (data) set(s => ({
          tokenUsage: { ...s.tokenUsage, [id]: {
            promptTokens: data.prompt_tokens || 0,
            completionTokens: data.completion_tokens || 0,
            totalTokens: data.total_tokens || 0,
          }}
        }))
      }).catch(() => {})
    },

    fetchRooms: async () => {
      const rooms = await fetch(`${API}/rooms`).then(r => r.json())
      set({ rooms: rooms || [] })
    },

    fetchAgents: async () => {
      const agents = await fetch(`${API}/agents`).then(r => r.json())
      set({ agents: agents || [] })
    },

    fetchSkills: async () => {
      const skills = await fetch(`${API}/skills`).then(r => r.json())
      set({ skills: skills || [] })
    },

    connectRoom: (roomId) => {
      const { ws } = get()
      // 已有连接且状态正常则跳过
      if (ws[roomId] && ws[roomId].readyState <= 1) return
      const socket = new WebSocket(`ws://${location.host}/ws/${roomId}`)

      // 断线自动重连
      socket.onclose = () => {
        set(s => {
          const newWs = { ...s.ws }
          delete newWs[roomId]
          return { ws: newWs }
        })
        // 3 秒后重连
        setTimeout(() => get().connectRoom(roomId), 3000)
      }

      // 心跳保活：每 30 秒发一次 ping
      const heartbeat = setInterval(() => {
        if (socket.readyState === WebSocket.OPEN) {
          socket.send(JSON.stringify({ type: 'ping' }))
        }
      }, 30000)
      socket.addEventListener('close', () => clearInterval(heartbeat))

      socket.onmessage = (e) => {
        const ev: WsEvent = JSON.parse(e.data)
        if (ev.type === 'agent_message') {
          set(s => {
            const msgs = [...(s.messages[roomId] || [])]

            // tool_use 消息：同一个 agent 只保留一条，持续更新状态
            if (ev.role === 'tool_use') {
              let idx = -1
              for (let i = msgs.length - 1; i >= 0; i--) {
                if (msgs[i].role === 'tool_use' && msgs[i].agent === ev.agent) { idx = i; break }
              }
              if (idx !== -1) {
                msgs[idx] = { ...msgs[idx], content: ev.content || msgs[idx].content, streaming: ev.streaming }
              } else if (ev.content) {
                msgs.push({ id: Date.now().toString(), agent: ev.agent, role: 'tool_use', content: ev.content, streaming: ev.streaming })
              }
              // 同步更新 todoItem 的 action 描述，同时将 assigned 转为 working
              const actionText = ev.streaming
                ? (ev.content === 'thinking' ? '正在整理结果...' : '正在搜索...')
                : undefined
              const todos = (s.todoItems[roomId] || []).map(t =>
                t.agent === ev.agent && (t.status === 'working' || t.status === 'assigned')
                  ? { ...t, status: 'working' as const, action: actionText || t.action }
                  : t
              )
              return {
                messages: { ...s.messages, [roomId]: msgs },
                todoItems: { ...s.todoItems, [roomId]: todos },
              }
            }

            // 普通消息：streaming 合并逻辑
            let idx = -1
            for (let i = msgs.length - 1; i >= 0; i--) {
              if (msgs[i].streaming && msgs[i].agent === ev.agent) { idx = i; break }
            }
            if (idx !== -1) {
              if (ev.action === 'replace') {
                // action=replace: 用新内容替换（文件驱动模式，文档已保存到 workspace）
                msgs[idx] = { ...msgs[idx], content: ev.content || '', streaming: false }
              } else {
                msgs[idx] = { ...msgs[idx], content: msgs[idx].content + (ev.content || ''), streaming: ev.streaming }
              }
            } else if (ev.action === 'replace') {
              // 找最后一条同 agent 消息替换，找不到则新增
              let found = false
              for (let i = msgs.length - 1; i >= 0; i--) {
                if (msgs[i].agent === ev.agent && msgs[i].role === ev.role) {
                  msgs[i] = { ...msgs[i], content: ev.content || '', streaming: false }
                  found = true
                  break
                }
              }
              if (!found && ev.content) {
                msgs.push({ id: Date.now().toString(), agent: ev.agent, role: ev.role, content: ev.content, streaming: false })
              }
            } else if (ev.content) {
              msgs.push({ id: Date.now().toString(), agent: ev.agent, role: ev.role, content: ev.content, streaming: ev.streaming, created_at: new Date().toISOString() })
            }
            return { messages: { ...s.messages, [roomId]: msgs } }
          })
        } else if (ev.type === 'room_status') {
          set(s => {
            let todos = [...(s.todoItems[roomId] || [])]
            if (ev.active_agent && ev.status === 'working') {
              // 更新活跃 agent 状态为 working，附带 action 描述
              todos = todos.map(t =>
                t.agent === ev.active_agent
                  ? { ...t, status: 'working' as const, action: ev.action || t.action }
                  : t
              )
            }
            // 记录工作开始时间（thinking/working 时设置，pending 时清除）
            const workStartedAt = { ...s.workStartedAt }
            if (ev.status === 'pending') {
              delete workStartedAt[roomId]
            } else if (!workStartedAt[roomId]) {
              workStartedAt[roomId] = Date.now()
            }
            // pending 时：清理所有残留 streaming 消息，只在所有任务都 done 后清空 todos
            if (ev.status === 'pending') {
              const cleanMsgs = (s.messages[roomId] || []).map(m =>
                m.streaming ? { ...m, streaming: false } : m
              )
              if (todos.length > 0) {
                const allDone = todos.every(t => t.status === 'done')
                if (allDone) todos = []
              }
              return {
                rooms: s.rooms.map(r => r.id === roomId
                  ? { ...r, status: ev.status, activeAgent: ev.active_agent, action: ev.action }
                  : r),
                messages: { ...s.messages, [roomId]: cleanMsgs },
                todoItems: { ...s.todoItems, [roomId]: todos },
                workingFiles: { ...s.workingFiles, [roomId]: [] },
                workStartedAt,
              }
            }
            return {
              rooms: s.rooms.map(r => r.id === roomId
                ? { ...r, status: ev.status, activeAgent: ev.active_agent, action: ev.action }
                : r),
              todoItems: { ...s.todoItems, [roomId]: todos },
              workStartedAt,
            }
          })
        } else if (ev.type === 'tasks_update') {
          set(s => ({ tasks: { ...s.tasks, [roomId]: ev.content } }))
        } else if (ev.type === 'workspace_file') {
          set(s => {
            const files = [...(s.workspace[roomId] || [])]
            if (!files.includes(ev.filename)) {
              files.push(ev.filename)
            }
            return { workspace: { ...s.workspace, [roomId]: files } }
          })
        } else if (ev.type === 'task_assign') {
          set(s => {
            const msgs = [...(s.messages[roomId] || [])]
            msgs.push({
              id: Date.now().toString(),
              agent: ev.from,
              role: 'task_assign' as const,
              content: ev.task,
              title: ev.to,
            })
            // 去重：同一 agent 只保留最新任务，清除已完成的
            const existing = (s.todoItems[roomId] || []).filter(t => t.status !== 'done' && t.agent !== ev.to)
            existing.push({
              id: `${Date.now()}-${ev.to}`,
              agent: ev.to,
              task: ev.task,
              status: 'assigned',
            })
            return {
              messages: { ...s.messages, [roomId]: msgs },
              todoItems: { ...s.todoItems, [roomId]: existing },
            }
          })
        } else if (ev.type === 'task_done') {
          set(s => {
            // 清理该 agent 的残留 streaming 消息
            const msgs = (s.messages[roomId] || []).map(m =>
              m.agent === ev.agent && m.streaming ? { ...m, streaming: false } : m
            )
            const todos = (s.todoItems[roomId] || []).map(t =>
              t.agent === ev.agent && t.status !== 'done'
                ? { ...t, status: 'done' as const }
                : t
            )
            return {
              messages: { ...s.messages, [roomId]: msgs },
              todoItems: { ...s.todoItems, [roomId]: todos },
            }
          })
        } else if (ev.type === 'mode_change') {
          set(s => ({
            rooms: s.rooms.map(r => r.id === roomId ? { ...r, mode: ev.mode } : r)
          }))
        } else if (ev.type === 'schedule_run') {
          set(s => {
            const msgs = [...(s.messages[roomId] || [])]
            msgs.push({
              id: Date.now().toString(),
              agent: ev.agent || 'scheduler',
              role: 'user' as const,
              content: `[定时任务] ${ev.content}`,
            })
            return { messages: { ...s.messages, [roomId]: msgs } }
          })
        } else if (ev.type === 'token_usage') {
          set(s => {
            const prev = s.tokenUsage[roomId] || { promptTokens: 0, completionTokens: 0, totalTokens: 0 }
            return {
              tokenUsage: {
                ...s.tokenUsage,
                [roomId]: {
                  promptTokens: prev.promptTokens + ev.prompt_tokens,
                  completionTokens: prev.completionTokens + ev.completion_tokens,
                  totalTokens: prev.totalTokens + ev.total_tokens,
                }
              }
            }
          })
        } else if (ev.type === 'artifact') {
          // 添加 artifact 消息卡片到聊天
          set(s => {
            const msgs = [...(s.messages[roomId] || [])]
            msgs.push({
              id: Date.now().toString(),
              agent: ev.agent,
              role: 'artifact' as const,
              content: ev.title || ev.filename,
              filename: ev.filename,
              title: ev.title,
            })
            // 同时更新 workspace 文件列表
            const files = [...(s.workspace[roomId] || [])]
            if (!files.includes(ev.filename)) {
              files.push(ev.filename)
            }
            return {
              messages: { ...s.messages, [roomId]: msgs },
              workspace: { ...s.workspace, [roomId]: files },
            }
          })
        } else if (ev.type === 'file_working') {
          set(s => {
            const files = [...(s.workingFiles[roomId] || [])]
            // 避免重复
            if (!files.some(f => f.agent === ev.agent && f.filename === ev.filename)) {
              files.push({ agent: ev.agent, filename: ev.filename, title: ev.title })
            }
            return { workingFiles: { ...s.workingFiles, [roomId]: files } }
          })
        } else if (ev.type === 'file_done') {
          set(s => {
            const files = (s.workingFiles[roomId] || []).filter(
              f => !(f.agent === ev.agent && f.filename === ev.filename)
            )
            return { workingFiles: { ...s.workingFiles, [roomId]: files } }
          })
        } else if (ev.type === 'file_read') {
          set(s => {
            const roomReaders = { ...(s.fileReaders[roomId] || {}) }
            const readers = [...(roomReaders[ev.filename] || [])]
            if (!readers.includes(ev.agent)) {
              readers.push(ev.agent)
            }
            roomReaders[ev.filename] = readers
            return { fileReaders: { ...s.fileReaders, [roomId]: roomReaders } }
          })
        }
      }
      set(s => ({ ws: { ...s.ws, [roomId]: socket } }))
    },

    sendMessage: (roomId, content) => {
      const { ws, user } = get()
      const userName = user?.name || '用户'
      const userMsg = { id: Date.now().toString(), agent: userName, role: 'user' as const, content, created_at: new Date().toISOString() }
      set(s => ({ messages: { ...s.messages, [roomId]: [...(s.messages[roomId] || []), userMsg] } }))
      ws[roomId]?.send(JSON.stringify({ type: 'user_message', content, user_name: userName }))
    },

    markFileRead: (roomId: string, filename: string) => {
      const userName = get().user?.name || '用户'
      set(s => {
        const roomReaders = { ...(s.fileReaders[roomId] || {}) }
        const readers = [...(roomReaders[filename] || [])]
        if (!readers.includes(userName)) {
          readers.push(userName)
        }
        roomReaders[filename] = readers
        return { fileReaders: { ...s.fileReaders, [roomId]: roomReaders } }
      })
    },

    setRoomTeam: async (roomId, team) => {
      const res = await fetch(`${API}/rooms/${roomId}/team`, {
        method: 'PUT', headers: { 'Content-Type': 'application/json' },
        body: JSON.stringify({ team })
      })
      if (res.ok) {
        const data = await res.json()
        set(s => ({
          rooms: s.rooms.map(r => r.id === roomId ? { ...r, team, master: data.master, members: data.members } : r)
        }))
      }
    },
  }
})