const WebSocket = require("ws")
const express = require("express")
const http = require("http")
const path = require("path")
const cors = require("cors")
const helmet = require("helmet")
const fs = require("fs")
const multer = require("multer")
const crypto = require("crypto")

const app = express()
const server = http.createServer(app)

// Ensure required directories exist
const dirs = ['data', 'public', 'public/emojis']
dirs.forEach(dir => {
  if (!fs.existsSync(dir)) fs.mkdirSync(dir, { recursive: true })
})

const uploadDirs = ['public/uploads']
uploadDirs.forEach(dir => {
  if (!fs.existsSync(dir)) fs.mkdirSync(dir, { recursive: true })
})

// Data file paths
const USERS_FILE = 'data/users.json'
const CHANNELS_FILE = 'data/channels.json'
const EMOJIS_FILE = 'data/emojis.json'

// Load or initialize data
let users = loadJSON(USERS_FILE, {})
let channels = loadJSON(CHANNELS_FILE, { 
  general: { 
    name: 'general', 
    isPublic: true, 
    messages: [],
    created: new Date().toISOString()
  }
})
let customEmojis = loadJSON(EMOJIS_FILE, {})

// Constants
const ADMIN_SIGNUP_KEY = "GqK25£1|p[3/"
const MAX_HISTORY = 100
const MAX_MESSAGE_LENGTH = 2000
const RATE_LIMIT = 50
const INACTIVE_DAYS = 3
const MAX_EMOJI_SIZE = 1024 * 1024 // 1MB

let activeUsers = new Map()

const userUpload = multer({
  storage: multer.diskStorage({
    destination: (req, file, cb) => cb(null, 'public/uploads/'),
    filename: (req, file, cb) => {
      const ext = path.extname(file.originalname)
      const name = crypto.randomBytes(8).toString('hex')
      cb(null, name + ext)
    }
  }),
  limits: { fileSize: 5 * 1024 * 1024 }, // 5MB max
  fileFilter: (req, file, cb) => {
    const allowed = ['image/png','image/jpeg','image/gif']
    cb(null, allowed.includes(file.mimetype))
  }
})

app.post("/api/upload-image", userUpload.single('image'), (req, res) => {
  if (!req.file) return res.status(400).json({ error: "No file" })
  res.json({ url: `/uploads/${req.file.filename}` })
})


// Security and CORS middleware
app.use(helmet({
  contentSecurityPolicy: false,
  crossOriginResourcePolicy: { policy: "cross-origin" }
}))
app.use(cors())
app.use(express.json())
app.use(express.static(path.join(__dirname, "public")))

// Configure multer for emoji uploads
const storage = multer.diskStorage({
  destination: (req, file, cb) => {
    cb(null, 'public/emojis/')
  },
  filename: (req, file, cb) => {
    const ext = path.extname(file.originalname)
    const name = crypto.randomBytes(8).toString('hex')
    cb(null, name + ext)
  }
})

const upload = multer({ 
  storage,
  limits: { fileSize: MAX_EMOJI_SIZE },
  fileFilter: (req, file, cb) => {
    const allowedTypes = ['image/png', 'image/gif']
    cb(null, allowedTypes.includes(file.mimetype))
  }
})

// Create WebSocket server
const wss = new WebSocket.Server({ server })

// Client tracking
let clients = new Map() // ws -> clientData
let channelClients = new Map() // channel -> Set of ws
let userRateLimits = new Map()

function loadJSON(filepath, defaultData) {
  try {
    if (fs.existsSync(filepath)) {
      return JSON.parse(fs.readFileSync(filepath, 'utf8'))
    }
  } catch (err) {
    console.error(`Error loading ${filepath}:`, err)
  }
  return defaultData
}

function saveJSON(filepath, data) {
  try {
    fs.writeFileSync(filepath, JSON.stringify(data, null, 2))
  } catch (err) {
    console.error(`Error saving ${filepath}:`, err)
  }
}

function generateId() {
  return crypto.randomBytes(6).toString('hex')
}

function getCurrentTimestamp() {
  return new Date().toISOString()
}

function sanitizeMessage(message) {
  // Don't sanitize HTML for admins - they can use HTML
  return message.trim().substring(0, MAX_MESSAGE_LENGTH)
}

function sanitizeUsername(username) {
  // Basic sanitization for usernames
  return username.replace(/<script[^>]*>.*?<\/script>/gi, '')
                .trim()
                .substring(0, 30)
}

function isRateLimited(username) {
  const now = Date.now()
  const userLimits = userRateLimits.get(username) || []
  const recentMessages = userLimits.filter(timestamp => now - timestamp < 60000)
  
  if (recentMessages.length >= RATE_LIMIT) return true
  
  recentMessages.push(now)
  userRateLimits.set(username, recentMessages)
  return false
}


function handleEditMessage(ws, msg) {
  const clientData = clients.get(ws)
  if (!clientData) return

  const { messageId, newContent } = msg
  const channelData = channels[clientData.currentChannel]
  
  if (!channelData) return

  const messageIndex = channelData.messages.findIndex(m => m.id === messageId)
  if (messageIndex === -1) {
    ws.send(JSON.stringify({
      type: "error",
      content: "Message not found"
    }))
    return
  }

  const message = channelData.messages[messageIndex]
  
  // Only allow editing own messages or admin can edit any
  if (message.from !== clientData.username && !clientData.isAdmin) {
    ws.send(JSON.stringify({
      type: "error",
      content: "Cannot edit this message"
    }))
    return
  }

  // Update message
  message.content = sanitizeMessage(newContent)
  message.edited = true
  message.editedAt = getCurrentTimestamp()
  
  saveJSON(CHANNELS_FILE, channels)
  
  broadcastToChannel(clientData.currentChannel, {
    type: "message_edited",
    messageId: messageId,
    newContent: message.content,
    edited: true,
    editedAt: message.editedAt
  })
}
function handleDeleteMessage(ws, msg) {
  const clientData = clients.get(ws)
  if (!clientData) return

  const { messageId } = msg
  const channelData = channels[clientData.currentChannel]
  
  if (!channelData) return

  const messageIndex = channelData.messages.findIndex(m => m.id === messageId)
  if (messageIndex === -1) {
    ws.send(JSON.stringify({
      type: "error",
      content: "Message not found"
    }))
    return
  }

  const message = channelData.messages[messageIndex]
  
  // Only allow deleting own messages or admin can delete any
  if (message.from !== clientData.username && !clientData.isAdmin) {
    ws.send(JSON.stringify({
      type: "error",
      content: "Cannot delete this message"
    }))
    return
  }

  // Remove message
  channelData.messages.splice(messageIndex, 1)
  saveJSON(CHANNELS_FILE, channels)
  
  broadcastToChannel(clientData.currentChannel, {
    type: "message_deleted",
    messageId: messageId
  })
}
function handleAdminBanUser(ws, msg) {
  const clientData = clients.get(ws)
  if (!clientData || !clientData.isAdmin) return

  const { username } = msg
  
  if (!users[username]) {
    ws.send(JSON.stringify({
      type: "error",
      content: "User not found"
    }))
    return
  }

  // Don't ban other admins
  if (users[username].isAdmin) {
    ws.send(JSON.stringify({
      type: "error", 
      content: "Cannot ban admin users"
    }))
    return
  }

  // Add banned flag
  users[username].banned = true
  users[username].bannedBy = clientData.username
  users[username].bannedAt = getCurrentTimestamp()
  saveJSON(USERS_FILE, users)

  // Disconnect the user if online
  for (let [clientWs, userData] of clients) {
    if (userData.username === username) {
      clientWs.close()
      break
    }
  }

  broadcast({
    type: "server",
    content: `${username} has been banned by ${clientData.username}`,
    timestamp: getCurrentTimestamp()
  })
}

function handleAdminMuteUser(ws, msg) {
  const clientData = clients.get(ws)
  if (!clientData || !clientData.isAdmin) return

  const { username, duration } = msg // duration in minutes
  
  if (!users[username]) {
    ws.send(JSON.stringify({
      type: "error",
      content: "User not found"
    }))
    return
  }

  const mutedUntil = Date.now() + (duration * 60 * 1000)
  users[username].mutedUntil = mutedUntil
  users[username].mutedBy = clientData.username
  saveJSON(USERS_FILE, users)

  broadcast({
    type: "server",
    content: `${username} has been muted for ${duration} minutes by ${clientData.username}`,
    timestamp: getCurrentTimestamp()
  })
}

function handleAdminClearChat(ws, msg) {
  const clientData = clients.get(ws)
  if (!clientData || !clientData.isAdmin) return

  const { channel } = msg
  const targetChannel = channel || clientData.currentChannel
  
  if (!channels[targetChannel]) return

  channels[targetChannel].messages = []
  saveJSON(CHANNELS_FILE, channels)

  broadcastToChannel(targetChannel, {
    type: "chat_cleared",
    clearedBy: clientData.username,
    timestamp: getCurrentTimestamp()
  })
}

function cleanupInactiveUsers() {
  const cutoff = Date.now() - (INACTIVE_DAYS * 24 * 60 * 60 * 1000)
  let cleaned = 0
  
  for (const [username, userData] of Object.entries(users)) {
    if (userData.lastSeen < cutoff && !userData.isAdmin) {
      delete users[username]
      cleaned++
    }
  }
  
  if (cleaned > 0) {
    saveJSON(USERS_FILE, users)
    console.log(`Cleaned up ${cleaned} inactive users`)
  }
}

// REST API endpoints
app.get("/", (req, res) => {
  res.sendFile(path.join(__dirname, "public", "index.html"))
})

app.get("/api/stats", (req, res) => {
  res.json({
    connectedUsers: clients.size,
    totalUsers: Object.keys(users).length,
    totalChannels: Object.keys(channels).length,
    publicChannels: Object.values(channels).filter(c => c.isPublic).length,
    uptime: process.uptime()
  })
})

app.get("/api/channels", (req, res) => {
  const { username } = req.query
  const userData = users[username]
  
  if (!userData) {
    return res.status(401).json({ error: "User not found" })
  }
  
  const channelList = Object.values(channels).map(channel => ({
    name: channel.name,
    isPublic: channel.isPublic,
    userCount: (channelClients.get(channel.name) || new Set()).size,
    canSee: channel.isPublic || userData.isAdmin
  })).filter(c => c.canSee)
  
  res.json(channelList)
})

app.get("/api/emojis", (req, res) => {
  const emojiDir = path.join(__dirname, "public/emojis")
  fs.readdir(emojiDir, (err, files) => {
    if (err) return res.status(500).json({ error: "Could not read emojis" })
    
    const emojis = files
      .filter(f => /\.(png|gif)$/i.test(f))
      .map(f => {
        const name = path.parse(f).name // filename without extension
        return {
          id: name,
          name: name,
          url: `/emojis/${f}`,
          usage: `:${name}:`
        }
      })
    
    res.json({ success: true, count: emojis.length, emojis })
  })
})


app.post("/api/signup", (req, res) => {
  const { username, password, adminKey } = req.body
  
  if (!username || !password) {
    return res.status(400).json({ error: "Username and password required" })
  }
  
  const cleanUsername = sanitizeUsername(username)
  if (cleanUsername.length < 2) {
    return res.status(400).json({ error: "Username must be at least 2 characters" })
  }
  
  if (users[cleanUsername]) {
    return res.status(409).json({ error: "Username already exists" })
  }
  
  const isAdmin = adminKey === ADMIN_SIGNUP_KEY
  
  users[cleanUsername] = {
    username: cleanUsername,
    password: password, // In production, hash this!
    isAdmin: isAdmin,
    color: '#93c5fd',
    created: getCurrentTimestamp(),
    lastSeen: Date.now()
  }
  
  saveJSON(USERS_FILE, users)
  
  res.json({ 
    success: true, 
    isAdmin: isAdmin,
    message: isAdmin ? "Admin account created" : "Account created"
  })
})

app.post("/api/login", (req, res) => {
  const { username, password } = req.body
  const userData = users[username]
  
  if (!userData || userData.password !== password) {
    return res.status(401).json({ error: "Invalid credentials" })
  }
  
  userData.lastSeen = Date.now()
  saveJSON(USERS_FILE, users)
  
  res.json({
    success: true,
    userData: {
      username: userData.username,
      isAdmin: userData.isAdmin,
      color: userData.color
    }
  })
})


app.post("/api/admin/execute", (req, res) => {
  const { username, code, scope, channel } = req.body
  const userData = users[username]
  
  if (!userData || !userData.isAdmin) {
    return res.status(401).json({ error: "Admin access required" })
  }
  
  const execution = {
    type: "admin_execute",
    code: code,
    scope: scope, // 'channel' or 'site'
    channel: channel,
    executor: username,
    timestamp: getCurrentTimestamp(),
    id: generateId()
  }
  
  if (scope === 'channel') {
    broadcastToChannel(channel, execution)
  } else {
    broadcast(execution)
  }
  
  res.json({ success: true })
})

// WebSocket connection handling
wss.on("connection", (ws, req) => {
  console.log(`New client connected from ${req.socket.remoteAddress}`)
  
  ws.on("message", (data) => {
    try {
      const msg = JSON.parse(data.toString())
      handleMessage(ws, msg)
    } catch (err) {
      console.error("Invalid message:", data.toString())
      ws.send(JSON.stringify({
        type: "error",
        content: "Invalid message format"
      }))
    }
  })

  ws.on("close", (code, reason) => {
    const clientData = clients.get(ws)
  if (clientData) {
    if (activeUsers.get(clientData.username) === ws) {
      activeUsers.delete(clientData.username)
    }
      
      // Remove from channel
      const channelSet = channelClients.get(clientData.currentChannel)
      if (channelSet) {
        channelSet.delete(ws)
      }
      
      const logoutMessage = {
        type: "logout",
        content: clientData.username,
        user_list: getChannelUsers(clientData.currentChannel),
        timestamp: getCurrentTimestamp()
      }
      
      broadcastToChannel(clientData.currentChannel, logoutMessage)
      console.log(`${clientData.username} disconnected from ${clientData.currentChannel}`)
    }
  })

  ws.on("error", (error) => {
    console.error("WebSocket error:", error)
  })
})

function handleMessage(ws, msg) {
  switch (msg.type) {
    case "auth":
      handleAuth(ws, msg)
      break
      
    case "join_channel":
      handleJoinChannel(ws, msg)
      break
      
    case "create_channel":
      handleCreateChannel(ws, msg)
      break
      
    case "user":
      handleUserMessage(ws, msg)
      break
      
    case "typing":
      handleTypingIndicator(ws, msg)
      break
      
    case "update_color":
      handleUpdateColor(ws, msg)
      break
      
    case "ping":
      ws.send(JSON.stringify({ type: "pong" }))
      break

      case "edit_message":
  handleEditMessage(ws, msg)
  break

case "delete_message":
  handleDeleteMessage(ws, msg)
  break

case "admin_ban_user":
  handleAdminBanUser(ws, msg)
  break

case "admin_mute_user":
  handleAdminMuteUser(ws, msg)
  break

case "admin_clear_chat":
  handleAdminClearChat(ws, msg)
  break

      
    default:
      ws.send(JSON.stringify({
        type: "error",
        content: "Unknown message type"
      }))
  }
}

function handleAuth(ws, msg) {
  const { username } = msg
  const userData = users[username]

  if (!userData) {
    ws.send(JSON.stringify({ type: "error", content: "User not authenticated" }))
    return
  }

  if (userData.banned) {
    ws.send(JSON.stringify({ type: "error", content: "You have been banned from this server" }))
    ws.close()
    return
  }

  // Kick previous connection if exists
  if (activeUsers.has(username)) {
    const oldWs = activeUsers.get(username)
    if (oldWs.readyState === WebSocket.OPEN) {
      oldWs.send(JSON.stringify({ type: "error", content: "Another person has signed in" }))
      oldWs.close()
    }
  }

  // mark this ws as active for username
  activeUsers.set(username, ws)

  userData.lastSeen = Date.now()

  const clientData = {
    username: username,
    isAdmin: userData.isAdmin,
    color: userData.color,
    currentChannel: 'general',
    isTyping: false
  }

  clients.set(ws, clientData)

  if (!channelClients.has('general')) channelClients.set('general', new Set())
  channelClients.get('general').add(ws)

  // send init data
  const channelData = channels['general']
  ws.send(JSON.stringify({
    type: "init",
    channel: 'general',
    history: channelData.messages.slice(-20),
    userData: { username, isAdmin: userData.isAdmin, color: userData.color },
    customEmojis: customEmojis
  }))

  broadcastToChannel('general', {
    type: "login",
    content: username,
    user_list: getChannelUsers('general'),
    timestamp: getCurrentTimestamp()
  })

  saveJSON(USERS_FILE, users)
}

function handleJoinChannel(ws, msg) {
  const clientData = clients.get(ws)
  if (!clientData) return
  
  const { channel } = msg
  const channelData = channels[channel]
  
  if (!channelData) {
    ws.send(JSON.stringify({
      type: "error",
      content: "Channel does not exist"
    }))
    return
  }
  
  if (!channelData.isPublic && !clientData.isAdmin) {
    ws.send(JSON.stringify({
      type: "error",
      content: "Access denied to private channel"
    }))
    return
  }
  
  // Leave current channel
  const oldChannelSet = channelClients.get(clientData.currentChannel)
  if (oldChannelSet) {
    oldChannelSet.delete(ws)
    broadcastToChannel(clientData.currentChannel, {
      type: "logout",
      content: clientData.username,
      user_list: getChannelUsers(clientData.currentChannel),
      timestamp: getCurrentTimestamp()
    })
  }
  
  // Join new channel
  clientData.currentChannel = channel
  if (!channelClients.has(channel)) {
    channelClients.set(channel, new Set())
  }
  channelClients.get(channel).add(ws)
  
  // Send channel data
  ws.send(JSON.stringify({
    type: "channel_switch",
    channel: channel,
    history: channelData.messages.slice(-20)
  }))
  
  // Announce join
  broadcastToChannel(channel, {
    type: "login",
    content: clientData.username,
    user_list: getChannelUsers(channel),
    timestamp: getCurrentTimestamp()
  })
}

function handleCreateChannel(ws, msg) {
  const clientData = clients.get(ws)
  if (!clientData) return
  
  const { channel, isPublic } = msg
  const channelName = channel.replace(/[^a-zA-Z0-9_-]/g, '').toLowerCase()
  
  if (!channelName || channelName.length < 2) {
    ws.send(JSON.stringify({
      type: "error",
      content: "Invalid channel name"
    }))
    return
  }
  
  if (channels[channelName]) {
    ws.send(JSON.stringify({
      type: "error",
      content: "Channel already exists"
    }))
    return
  }
  
  channels[channelName] = {
    name: channelName,
    isPublic: !!isPublic,
    messages: [],
    created: getCurrentTimestamp(),
    createdBy: clientData.username
  }
  
  saveJSON(CHANNELS_FILE, channels)
  
  ws.send(JSON.stringify({
    type: "channel_created",
    channel: channelName
  }))
}

function handleUserMessage(ws, msg) {
  const clientData = clients.get(ws)
  if (!clientData) return

  // Check if user is muted
  const userData = users[clientData.username]
  if (userData && userData.mutedUntil && userData.mutedUntil > Date.now()) {
    ws.send(JSON.stringify({
      type: "error",
      content: "You are muted"
    }))
    return
  }

  if (isRateLimited(clientData.username)) {
    ws.send(JSON.stringify({
      type: "error",
      content: "Rate limit exceeded"
    }))
    return
  }

  let content = msg.content.trim()
  if (!content) return

  // Process emojis
  content = content.replace(/:([a-zA-Z0-9_]+):/g, (match, emojiName) => {
    if (customEmojis[emojiName]) {
      return `<img class="emoji" src="/emojis/${customEmojis[emojiName].filename}" alt=":${emojiName}:" title=":${emojiName}:">`
    }
    return match
  })

  // Sanitize content for non-admins
  if (!clientData.isAdmin) {
    content = content.replace(/<script[^>]*>.*?<\/script>/gi, '')
                    .replace(/<[^>]*>/g, '')
  }

  const messageData = {
    type: "user",
    from: clientData.username,
    content: content,
    color: clientData.color,
    isAdmin: clientData.isAdmin,
    timestamp: getCurrentTimestamp(),
    id: generateId()
  }

  // Add to channel history
  const channelData = channels[clientData.currentChannel]
  if (channelData) {
    channelData.messages.push(messageData)
    if (channelData.messages.length > MAX_HISTORY) {
      channelData.messages = channelData.messages.slice(-MAX_HISTORY)
    }
    saveJSON(CHANNELS_FILE, channels)
  }

  broadcastToChannel(clientData.currentChannel, messageData)
}

function handleTypingIndicator(ws, msg) {
  const clientData = clients.get(ws)
  if (!clientData) return
  
  clientData.isTyping = msg.isTyping
  
  const typingUsers = getChannelUsers(clientData.currentChannel)
    .filter(user => {
      for (let [clientWs, data] of clients) {
        if (data.username === user && data.isTyping && data.currentChannel === clientData.currentChannel) {
          return true
        }
      }
      return false
    })
  
  broadcastToChannel(clientData.currentChannel, {
    type: "typing",
    users: typingUsers
  }, ws)
}

function handleUpdateColor(ws, msg) {
  const clientData = clients.get(ws)
  if (!clientData) return
  
  const { color } = msg
  if (!/^#[0-9a-fA-F]{6}$/.test(color)) {
    ws.send(JSON.stringify({
      type: "error",
      content: "Invalid color format"
    }))
    return
  }
  
  clientData.color = color
  users[clientData.username].color = color
  saveJSON(USERS_FILE, users)
  
  ws.send(JSON.stringify({
    type: "color_updated",
    color: color
  }))
}

function getChannelUsers(channel) {
  const channelSet = channelClients.get(channel)
  if (!channelSet) return []
  
  return Array.from(channelSet).map(ws => {
    const clientData = clients.get(ws)
    return clientData ? clientData.username : null
  }).filter(Boolean)
}

function broadcastToChannel(channel, data, excludeWs = null) {
  const channelSet = channelClients.get(channel)
  if (!channelSet) return
  
  const str = JSON.stringify(data)
  for (let client of channelSet) {
    if (client !== excludeWs && client.readyState === WebSocket.OPEN) {
      try {
        client.send(str)
      } catch (err) {
        console.error("Error sending to client:", err)
      }
    }
  }
}

function broadcast(data, excludeWs = null) {
  const str = JSON.stringify(data)
  for (let client of wss.clients) {
    if (client !== excludeWs && client.readyState === WebSocket.OPEN) {
      try {
        client.send(str)
      } catch (err) {
        console.error("Error sending to client:", err)
      }
    }
  }
}

// Cleanup inactive users every hour
setInterval(cleanupInactiveUsers, 60 * 60 * 1000)

// Run cleanup every 5 minutes for other maintenance
setInterval(() => {
  const now = Date.now()
  for (let [username, timestamps] of userRateLimits) {
    const recent = timestamps.filter(t => now - t < 60000)
    if (recent.length === 0) {
      userRateLimits.delete(username)
    } else {
      userRateLimits.set(username, recent)
    }
  }
}, 5 * 60 * 1000)

// Graceful shutdown
process.on('SIGTERM', () => {
  console.log('Shutting down gracefully...')
  broadcast({
    type: "server",
    content: "Server shutting down for maintenance",
    timestamp: getCurrentTimestamp()
  })
  
  setTimeout(() => {
    server.close(() => {
      console.log('Server closed')
      process.exit(0)
    })
  }, 1000)
})

const PORT = process.env.PORT || 1242
server.listen(PORT, () => {
  console.log(`Server running on http://localhost:${PORT}`)
  console.log(`Admin signup key: ${ADMIN_SIGNUP_KEY}`)
})