import { create } from "zustand"
import type { Agent, AppState, GatewayProvider, RoleMapping, LedgerEntry, AgentRoleConfig, RuntimeRole, MCPConfig } from "@/types"
import { generateStableProviderId } from "@/lib/migration"
import { normalizeRole } from "@/lib/role-identity"
import { ALL_ROLES, type RoleDefinition } from "@/runtime/runtime-role-registry"
import { RuntimeOS } from "@/runtime/RuntimeOS"
import { setApiKey, removeApiKey } from "@/lib/secure-storage"

export interface ThinkingConfig {
  budget: number
  visualizationMode: "rainbow" | "classic"
  collapseBehavior: "auto" | "always_expanded" | "always_collapsed"
  showTokenCount: boolean
  showElapsedTime: boolean
}

const LOG_PREFIX = "[AppStore]"
const PLAN_MODE_KEY = 'agentic-plan-mode'
const SANDBOX_MODE_KEY = 'agentic-sandbox-mode'
const THINKING_CONFIG_KEY = 'agentic-thinking-config'
const OUTPUT_STYLE_KEY = 'agentic-output-style'

const DEFAULT_THINKING_CONFIG: ThinkingConfig = {
  budget: 0,
  visualizationMode: "rainbow",
  collapseBehavior: "auto",
  showTokenCount: true,
  showElapsedTime: true,
}

function loadPlanMode(): 'auto' | 'always' | 'never' {
  try {
    const stored = localStorage.getItem(PLAN_MODE_KEY)
    if (stored === 'auto' || stored === 'always' || stored === 'never') return stored
  } catch { console.warn("[Store] Failed to load plan mode from localStorage") }
  return 'auto'
}

function persistPlanMode(mode: 'auto' | 'always' | 'never'): void {
  try {
    localStorage.setItem(PLAN_MODE_KEY, mode)
  } catch { console.warn("[Store] Failed to persist plan mode") }
}

function loadSandboxMode(): 'on' | 'off' {
  try {
    const stored = localStorage.getItem(SANDBOX_MODE_KEY)
    if (stored === 'on' || stored === 'off') return stored
  } catch { console.warn("[Store] Failed to load sandbox mode") }
  return 'on'
}

function persistSandboxMode(mode: 'on' | 'off'): void {
  try {
    localStorage.setItem(SANDBOX_MODE_KEY, mode)
  } catch { console.warn("[Store] Failed to persist sandbox mode") }
}

function loadThinkingConfig(): ThinkingConfig {
  try {
    const stored = localStorage.getItem(THINKING_CONFIG_KEY)
    if (stored) {
      const parsed = JSON.parse(stored)
      return { ...DEFAULT_THINKING_CONFIG, ...parsed }
    }
  } catch { console.warn("[Store] Failed to load thinking config") }
  return { ...DEFAULT_THINKING_CONFIG }
}

function persistThinkingConfig(config: ThinkingConfig): void {
  try {
    localStorage.setItem(THINKING_CONFIG_KEY, JSON.stringify(config))
  } catch { console.warn("[Store] Failed to persist thinking config") }
}

export type OutputStyle = "default" | "explanatory" | "learning"

function loadOutputStyle(): OutputStyle {
  try {
    const stored = localStorage.getItem(OUTPUT_STYLE_KEY)
    if (stored === "default" || stored === "explanatory" || stored === "learning") return stored
  } catch { console.warn("[Store] Failed to load output style") }
  return "default"
}

function persistOutputStyle(style: OutputStyle): void {
  try {
    localStorage.setItem(OUTPUT_STYLE_KEY, style)
  } catch { console.warn("[Store] Failed to persist output style") }
}

function log(...args: unknown[]) {
  console.log(LOG_PREFIX, ...args)
}

function warn(...args: unknown[]) {
  console.warn(LOG_PREFIX, "[WARN]", ...args)
}

function defToRoleConfig(def: RoleDefinition): AgentRoleConfig {
  return {
    id: def.id,
    name: def.name,
    runtimeRole: def.runtimeRole,
    description: def.description,
    color: def.color,
    icon: def.icon,
    temperature: def.temperature,
    maxTokens: def.maxTokens,
    systemPrompt: def.systemPrompt,
    systemPromptVersion: 1,
    runtimeState: "idle",
    capabilities: { ...def.capabilities },
    toolPermissions: [...def.toolPermissions],
    memoryScope: def.memoryScope,
    verificationRequired: false,
    priority: def.priority,
    collaborationTags: [...def.collaborationTags],
    isBuiltIn: true,
    isEnabled: true,
    executionCount: 0,
  }
}

function getDefaultRoles(): AgentRoleConfig[] {
  return ALL_ROLES.map(defToRoleConfig)
}

interface AppStore {
  appState: AppState
  agents: Agent[]
  providers: GatewayProvider[]
  roleConfigs: AgentRoleConfig[]
  roleMappings: RoleMapping[]
  ledger: LedgerEntry[]
  mcpServers: MCPConfig[]
  defaultsInitialized: boolean
  planMode: "auto" | "always" | "never"
  setPlanMode: (mode: "auto" | "always" | "never") => void
  sandboxMode: "on" | "off"
  setSandboxMode: (mode: "on" | "off") => void
  thinkingConfig: ThinkingConfig
  setThinkingConfig: (config: ThinkingConfig) => void
  outputStyle: OutputStyle
  setOutputStyle: (style: OutputStyle) => void
  debugMode: boolean
  setDebugMode: (enabled: boolean) => void
  mockMode: boolean
  setMockMode: (enabled: boolean) => void
  setAppState: (state: AppState) => void
  updateAgent: (agentId: string, updates: Partial<Agent>) => void
  addProvider: (provider: GatewayProvider) => void
  updateProvider: (providerId: string, updates: Partial<GatewayProvider>) => void
  removeProvider: (providerId: string) => void
  upsertRoleConfig: (config: AgentRoleConfig) => void
  removeRoleConfig: (id: string) => void
  setRoleMapping: (mapping: RoleMapping) => void
  addLedgerEntry: (entry: LedgerEntry) => void
  resetAllAgentUsage: () => void
  initializeDefaultRoles: () => void
  getAllModels: () => { providerId: string; providerName: string; models: import("@/types").ProviderModel[] }[]
  addMcpServer: (config: MCPConfig) => void
  removeMcpServer: (id: string) => void
  toggleMcpServer: (id: string) => void
  updateMcpServer: (id: string, updates: Partial<MCPConfig>) => void
}

function findDependentRoles(
  roleConfigs: AgentRoleConfig[],
  providerId: string,
): AgentRoleConfig[] {
  return roleConfigs.filter((r) => r.providerId === providerId)
}

export const useAppStore = create<AppStore>((set, get) => ({
  appState: "idle",
  agents: [
    { id: "agent-coding", name: "Coding Agent", role: "coder", state: "idle", currentTask: null, tokenUsage: 0, model: "" },
    { id: "agent-design", name: "Design Agent", role: "design", state: "idle", currentTask: null, tokenUsage: 0, model: "" },
    { id: "agent-vision", name: "Vision Agent", role: "vision", state: "idle", currentTask: null, tokenUsage: 0, model: "" },
    { id: "agent-qa", name: "QA Agent", role: "qa", state: "idle", currentTask: null, tokenUsage: 0, model: "" },
    { id: "agent-manager", name: "Manager Agent", role: "manager", state: "idle", currentTask: null, tokenUsage: 0, model: "" },
    { id: "agent-runtime", name: "Runtime Agent", role: "runtime", state: "idle", currentTask: null, tokenUsage: 0, model: "" },
  ],
  providers: [],
  roleConfigs: [],
  roleMappings: [],
  ledger: [],
  mcpServers: [],
  defaultsInitialized: false,
  planMode: loadPlanMode(),
  setPlanMode: (mode) => {
    persistPlanMode(mode)
    set({ planMode: mode })
  },
  sandboxMode: loadSandboxMode(),
  setSandboxMode: (mode) => {
    persistSandboxMode(mode)
    set({ sandboxMode: mode })
  },
  thinkingConfig: loadThinkingConfig(),
  setThinkingConfig: (config) => {
    persistThinkingConfig(config)
    set({ thinkingConfig: config })
  },
  outputStyle: loadOutputStyle(),
  setOutputStyle: (style) => {
    persistOutputStyle(style)
    set({ outputStyle: style })
  },
  debugMode: false,
  setDebugMode: (enabled) => set({ debugMode: enabled }),
  mockMode: true,
  setMockMode: (enabled) => set({ mockMode: enabled }),
  setAppState: (state) => set({ appState: state }),
  updateAgent: (agentId, updates) =>
    set((store) => ({
      agents: store.agents.map((a) => (a.id === agentId ? { ...a, ...updates } : a)),
    })),
  addProvider: (provider) => {
    const existingIds = get().providers.map((p) => p.id)
    const stableId = generateStableProviderId(provider, existingIds)

    set((store) => {
      if (existingIds.includes(stableId)) {
        warn(`addProvider: duplicate slug "${stableId}" for provider "${provider.name}" — generating unique id`)
      }

      const normalized: GatewayProvider = {
        ...provider,
        id: stableId,
      }

      log(`addProvider: "${normalized.name}" (id: ${normalized.id})`)
      const providers = [...store.providers, normalized]
      const roleConfigs = store.roleConfigs.map((role) => {
        if (!role.isEnabled) return role
        if (role.providerId && role.model) return role
        return {
          ...role,
          providerId: role.providerId ?? normalized.id,
          model: role.model ?? normalized.models[0]?.id,
        }
      })
      return { providers, roleConfigs }
    })

    if (provider.apiKey) {
      setApiKey(stableId, provider.apiKey).catch(() => {})
    }
    // Auto-disable mock mode when a real provider is added
    if (get().mockMode) {
      set({ mockMode: false })
    }
  },
  updateProvider: (providerId, updates) => {
    set((store) => {
      if (!store.providers.some((p) => p.id === providerId)) {
        warn(`updateProvider: provider "${providerId}" not found`)
        return store
      }
      log(`updateProvider: "${providerId}"`)
      return { providers: store.providers.map((p) => (p.id === providerId ? { ...p, ...updates } : p)) }
    })
    if ("apiKey" in updates && updates.apiKey) {
      setApiKey(providerId, updates.apiKey as string).catch(() => {})
    }
  },
  removeProvider: (providerId) => {
    set((store) => {
      const dependentRoles = findDependentRoles(store.roleConfigs, providerId)
      if (dependentRoles.length > 0) {
        warn(`removeProvider: "${providerId}" has ${dependentRoles.length} dependent role(s): ${dependentRoles.map((r) => r.name).join(", ")}`)
        log(`removeProvider: clearing providerId from ${dependentRoles.length} role(s)`)
        return {
          providers: store.providers.filter((p) => p.id !== providerId),
          roleConfigs: store.roleConfigs.map((r) =>
            r.providerId === providerId ? { ...r, providerId: undefined, model: undefined } : r,
          ),
        }
      }
      log(`removeProvider: "${providerId}" removed (no dependents)`)
      return { providers: store.providers.filter((p) => p.id !== providerId) }
    })
    removeApiKey(providerId).catch(() => {})
  },
  upsertRoleConfig: (config) =>
    set((store) => ({
      roleConfigs: store.roleConfigs.some((r) => r.id === config.id)
        ? store.roleConfigs.map((r) => (r.id === config.id ? config : r))
        : [...store.roleConfigs, config],
    })),
  removeRoleConfig: (id) =>
    set((store) => ({ roleConfigs: store.roleConfigs.filter((r) => r.id !== id) })),
  setRoleMapping: (mapping) =>
    set((store) => ({
      roleMappings: [
        ...store.roleMappings.filter((r) => r.role !== mapping.role),
        mapping,
      ],
    })),

  /** Keep ledger capped at 500 entries (newest) */
  addLedgerEntry: (entry) =>
    set((store) => {
      const ledger = [...store.ledger, entry]
      if (ledger.length > 500) ledger.splice(0, ledger.length - 500)
      return { ledger }
    }),
  resetAllAgentUsage: () =>
    set((store) => ({
      agents: store.agents.map((a) => ({ ...a, tokenUsage: 0 })),
    })),
  initializeDefaultRoles: () =>
    set((store) => {
      if (store.defaultsInitialized) return store
      if (store.roleConfigs.length > 0) {
        const defaults = getDefaultRoles()
        const backfilled = store.roleConfigs.map((r) => ({
          ...r,
          runtimeRole: r.runtimeRole ?? defaults.find((d) => d.id === r.id)?.runtimeRole ?? normalizeRole(r.id) ?? r.id as RuntimeRole,
        }))
        return { roleConfigs: backfilled, defaultsInitialized: true }
      }
      return {
        roleConfigs: getDefaultRoles(),
        defaultsInitialized: true,
      }
    }),
  getAllModels: () => {
    const providers = get().providers
    return providers
      .filter((p) => p.models.length > 0)
      .map((p) => ({
        providerId: p.id,
        providerName: p.name,
        models: p.models,
      }))
  },
  addMcpServer: (config) =>
    set((store) => {
      const existing = store.mcpServers.some((m) => m.id === config.id)
      if (existing) return store
      const servers = [...store.mcpServers, config]
      try {
        const runtime = RuntimeOS.getInstance()
        runtime.mcpServerManager.addServer({
          name: config.name,
          transport: { type: "stdio", command: config.command, args: config.args },
          enabled: config.enabled,
        })
      } catch (err) {
        console.warn("[app-store] Failed to add MCP server to RuntimeOS:", err)
      }
      return { mcpServers: servers }
    }),
  removeMcpServer: (id) =>
    set((store) => {
      const server = store.mcpServers.find((m) => m.id === id)
      if (server) {
        try {
          const runtime = RuntimeOS.getInstance()
          runtime.mcpServerManager.removeServer(server.name)
        } catch (err) {
          console.warn("[app-store] Failed to remove MCP server from RuntimeOS:", err)
        }
      }
      return { mcpServers: store.mcpServers.filter((m) => m.id !== id) }
    }),
  toggleMcpServer: (id) =>
    set((store) => ({
      mcpServers: store.mcpServers.map((m) =>
        m.id === id ? { ...m, enabled: !m.enabled } : m,
      ),
    })),
  updateMcpServer: (id, updates) =>
    set((store) => ({
      mcpServers: store.mcpServers.map((m) =>
        m.id === id ? { ...m, ...updates } : m,
      ),
    })),
}))
