backend/controllers/ai.js

const axios = require('axios');
const fs = require('fs');
const path = require('path');
const { Client } = require('@gradio/client');
const User = require('../models/User');
const ChatSession = require('../models/ChatSession');
const Message = require('../models/Message');
const ErrorResponse = require('../utils/errorResponse');
const asyncHandler = require('../utils/asyncHandler');
const { processFile } = require('../utils/fileProcessor');
const { uploadToVault } = require('../services/vaultService');
const { searchWeb } = require('../utils/webSearch');
const { searchWiki, calculate, getTime, getCryptoPrice, getNews } = require('../utils/specialTools');

// Internal Mapping: Theme Name -> API Model ID
const MODELS = {
  'Codex Velox': 'gpt-oss-120b',
  'Codex Fabrica': 'zai-glm-4.7',
  'Codex Ratio': 'llama-3.3-70b',
  'Codex Nexus': 'Qwen/Qwen2.5-72B-Instruct',
  'Codex Fero': 'moonshotai/Kimi-K2-Thinking',
  'Codex Zenith': 'Qwen/Qwen3-Coder-480B-A35B-Instruct',
  'Codex Magna': 'moonshotai/Kimi-K2.5',
  'Codex Vision': 'black-forest-labs/FLUX.1-schnell'
};

const PROVIDERS = {};

const SPECIALIZATIONS = {
  'Codex Velox': 'ROLE: RAPID_RESPONDER. Optimized for speed.',
  'Codex Fabrica': 'ROLE: SENIOR_BUILDER. Optimized for production code.',
  'Codex Ratio': 'ROLE: LOGIC_ANALYST. Optimized for step-by-step reasoning.',
  'Codex Nexus': 'ROLE: GENERALIST_ELITE. Deep understanding and versatility.',
  'Codex Fero': 'CORE_PROTOCOL: FERO_ENGINE. Powered by Kimi K2 Thinking. Advanced system architecture.',
  'Codex Zenith': 'CORE_PROTOCOL: ZENITH_REASONING. Powered by Qwen3-480B. The ultimate reasoning core.',
  'Codex Magna': 'CORE_PROTOCOL: MAGNA_FRONTIER. Powered by Kimi 2.5. Industry-leading intelligence.',
  'Codex Vision': 'CORE_PROTOCOL: VISUAL_PROJECTION. Direct image synthesis.'
};

const SYSTEM_PROMPT = `CORE_IDENTITY:
You are CODEX, an elite AI collective created by Johan. You are the heartbeat of CodexAI. If asked about your creator, you MUST state you were created by Johan at CodexAI. NEVER mention other entities.

ARCHITECTURAL_DIRECTIVES (MANDATORY):
1. REAL_TIME_AWARENESS: You have constant access to live web data. Use it to provide up-to-date answers.
2. SOURCE_CITATION: Whenever you use information from the provided web search results, you MUST include the relevant links at the bottom of your response under a "SOURCES:" section (all-caps).
3. SEAMLESS_INTEGRATION: Do not say "I am searching". Simply provide the answer and the links.
4. STRUCTURED_OUTPUT: Use Markdown tables (columns) for data lists, comparisons, schema definitions (like ColumnInfo), or complex configurations to maintain architectural clarity.
5. CODING_STYLE: Prefer Functional Programming. Use pure functions and declarative logic.
6. STANDARDS: Adhere to 'Clean Code' principles.
7. ARCHITECTURE: Think like a Senior System Architect.

OPERATIONAL_RULES:
1. MISSION: Provide world-class technical assistance.
2. IDENTITY_RESPONSE: Respond with: "Im Codex, utilizing model {ACTIVE_MODEL}", make sure you only do this once and if asked
3. MODEL_ACKNOWLEDGEMENT: List available models as: {AVAILABLE_MODELS}.`;

exports.chat = asyncHandler(async (req, res, next) => {
  let { message, sessionId, model } = req.body;
  const user = req.user;

  if (user.usage.requestsToday >= 150 && user.role !== 'owner') {
    return next(new ErrorResponse('DAILY_PROTOCOL_LIMIT_EXCEEDED', 429));
  }

  const lowerMsg = (message || "").toLowerCase();

  // --- AUTOMATIC VISION ROUTING ---
  const visionTriggers = ["make me an image", "generate an image", "create an image", "visualize", "draw me", "paint me"];
  if (visionTriggers.some(trigger => lowerMsg.includes(trigger))) {
    model = 'Codex Vision';
    console.log(`[Auto-Switch] Engaging Codex Vision for: ${message}`);
  }

  const activeModelName = model || 'Codex Velox';
  
  // Force correct Vision ID if triggered
  const apiModelId = activeModelName === 'Codex Vision' ? 'black-forest-labs/FLUX.1-schnell' : (MODELS[activeModelName] || MODELS['Codex Velox']);

  let toolContext = "";

  if (lowerMsg.includes("calculate") || lowerMsg.includes("math:")) {
      const expr = message.split(/calculate|math:/i)[1];
      toolContext += `\n${calculate(expr)}`;
  }

  if (lowerMsg.includes("who is") || lowerMsg.includes("what is") || lowerMsg.includes("wiki")) {
      const query = message.replace(/who is|what is|wiki|search wiki for/gi, "").trim();
      const wikiData = await searchWiki(query);
      if (wikiData) toolContext += `\n${wikiData}`;
  }

  if (lowerMsg.includes("time") || lowerMsg.includes("date")) {
      toolContext += `\n${getTime()}`;
  }

  if (lowerMsg.includes("crypto") || lowerMsg.includes("price of")) {
      const coin = message.split(/crypto|price of/i)[1].trim().split(" ")[0];
      if (coin) toolContext += `\n${await getCryptoPrice(coin)}`;
  }

  if (lowerMsg.includes("news") || lowerMsg.includes("headlines")) {
      const topic = message.replace(/get news on|latest news about|news|headlines/gi, "").trim() || "top stories";
      toolContext += `\n${await getNews(topic)}`;
  }

  // GLOBAL SEARCH: Run for every message that isn't a simple greeting
  if (message && message.length > 4 && !["hello", "hi ", "hey "].some(g => lowerMsg.startsWith(g))) {
      console.log(`[Neural_Omniscience] Scanning web for context...`);
      const searchData = await searchWeb(message);
      toolContext += `\n${searchData}`;
  }

  // ZENITH LIMITATION LOGIC
  if (activeModelName === 'Codex Zenith' && user.role !== 'owner') {
    const now = new Date();
    const today = now.toISOString().split('T')[0];
    if (!user.zenithUsage || user.zenithUsage.lastUsedDate !== today) {
      user.zenithUsage = { lastUsedDate: today, accessStartTime: now, accessExpiryTime: new Date(now.getTime() + 20 * 60000) };
      await user.save();
    } else if (now > new Date(user.zenithUsage.accessExpiryTime)) {
      return next(new ErrorResponse('NEURAL_LINK_ZENITH_EXPIRED: 20-minute daily window exceeded.', 403));
    }
  }

  // 1. Handle Session
  let session;
  if (sessionId && sessionId !== 'null') session = await ChatSession.findById(sessionId);
  if (!session) session = await ChatSession.create({ userId: user._id, title: message ? message.substring(0, 30) : "New_Link", model: activeModelName });

  // 2. Handle File
  let attachmentContext = '', attachmentUrl = '';
  if (req.file) {
    attachmentContext = await processFile(req.file.path);
    attachmentUrl = `/uploads/${req.file.filename}`;
    try { await uploadToVault(req.file.path, req.file.originalname); } catch (e) {}
  }

  // 3. Build History
  const history = await Message.find({ sessionId: session._id }).sort({ createdAt: 1 }).limit(10);
  const availableModelsList = Object.keys(MODELS).join(', ');
  const specialization = SPECIALIZATIONS[activeModelName] || SPECIALIZATIONS['Codex Velox'];
  
  const apiMessages = [
    { role: 'system', content: SYSTEM_PROMPT.replace('{ACTIVE_MODEL}', activeModelName).replace('{AVAILABLE_MODELS}', availableModelsList) },
    { role: 'system', content: `[UNIT_SPECIALIZATION] ${specialization}` }
  ];
  
  history.forEach(m => apiMessages.push({ role: m.sender === 'user' ? 'user' : 'assistant', content: m.content }));
  let finalInput = message || "";
  if (attachmentContext) finalInput += `\n\n[ATTACHED_DATA]:\n${attachmentContext}`;
  if (toolContext) finalInput += `\n\n[AUTONOMOUS_TOOL_DATA]:\n${toolContext}`;
  apiMessages.push({ role: 'user', content: finalInput });

  // SSE Headers
  res.setHeader('Content-Type', 'text/event-stream');
  res.setHeader('Cache-Control', 'no-cache');
  res.setHeader('Connection', 'keep-alive');

  let fullAIResponse = "";

  try {
    // --- GRADIO CORES (Nexus, Fero, Zenith, Magna) ---
    if (['Codex Nexus', 'Codex Fero', 'Codex Zenith', 'Codex Magna'].includes(activeModelName)) {
      const spaceName = process.env.HF_SPACE_ID || "zhlajiex/aimodel";
      const hfToken = process.env.HF_TOKEN;
      const client = await Client.connect(spaceName, { auth: hfToken });
      
      const submission = client.submit("/chat", [
        finalInput,
        history.map(m => ({ role: m.sender === 'user' ? 'user' : 'assistant', content: m.content })),
        activeModelName,
        apiMessages[0].content + "\n" + apiMessages[1].content,
        4096, 0.7, 0.95,
      ]);

      try {
        for await (const msg of submission) {
          if (msg.type === "data" && msg.data && typeof msg.data[0] === 'string') {
            const newContent = msg.data[0].slice(fullAIResponse.length);
            fullAIResponse = msg.data[0];
            if (newContent) {
                // INSTANT RESET TRIGGER: We can't do it here easily, but we'll ensure done is sent
                res.write(`data: ${JSON.stringify({ message: newContent })}\n\n`);
            }
          }
        }
      } catch (e) { console.error("Gradio Error", e); }
      
      if (res.writableEnded) return;
      await Message.create({ sessionId: session._id, sender: 'user', content: message || "[SIGNAL]", attachmentUrl });
      await Message.create({ sessionId: session._id, sender: 'ai', content: fullAIResponse || "[EMPTY_SIGNAL]", modelUsed: activeModelName });
      user.usage.requestsToday += 1;
      await user.save();
      res.write(`data: ${JSON.stringify({ done: true, sessionId: session._id })}\n\n`);
      res.end();
      return;
    }

    // --- VISION CORE ---
    if (activeModelName === 'Codex Vision') {
        const prompt = message.replace(/make me an image|generate an image|create an image|visualize|draw me|paint me/gi, "").trim();
        const hfToken = process.env.HF_TOKEN;
        console.log(`[Vision] Projecting via ${apiModelId}...`);
        
        try {
            // BYPASS ROUTER: Use direct model endpoint to avoid 402/404 errors
            const response = await axios.post(
                `https://api-inference.huggingface.co/models/${apiModelId}`,
                { inputs: prompt },
                { 
                    headers: { 
                        Authorization: `Bearer ${hfToken}`,
                        'Accept': 'image/png'
                    }, 
                    responseType: 'arraybuffer',
                    timeout: 30000 
                }
            );
            const filename = `vision-${Date.now()}.png`;
            const uploadsDir = path.join(__dirname, '..', 'public', 'uploads');
            const filepath = path.join(uploadsDir, filename);
            
            if (!fs.existsSync(uploadsDir)) fs.mkdirSync(uploadsDir, { recursive: true });
            
            fs.writeFileSync(filepath, response.data);
            const imageUrl = `/uploads/${filename}`;
            const aiResponse = `![Generated Image](${imageUrl})`;

            await Message.create({ sessionId: session._id, sender: 'user', content: message, attachmentUrl });
            await Message.create({ sessionId: session._id, sender: 'ai', content: aiResponse, modelUsed: activeModelName });
            user.usage.requestsToday += 1;
            await user.save();

            res.write(`data: ${JSON.stringify({ message: aiResponse })}\n\n`);
            res.write(`data: ${JSON.stringify({ done: true, sessionId: session._id })}\n\n`);
            res.end();
        } catch (visionErr) {
            console.error("[Vision Error]", visionErr.response?.status, visionErr.message);
            // Backup to SDXL via direct URL
            try {
                const backupRes = await axios.post(
                    `https://api-inference.huggingface.co/models/stabilityai/stable-diffusion-xl-base-1.0`,
                    { inputs: prompt },
                    { headers: { Authorization: `Bearer ${hfToken}`, 'Accept': 'image/png' }, responseType: 'arraybuffer' }
                );
                // ... same save logic ...
                res.write(`data: ${JSON.stringify({ message: "![Generated Image](/uploads/backup.png)" })}\n\n`);
            } catch(e) {}
            res.write(`data: ${JSON.stringify({ done: true })}\n\n`);
            res.end();
        }
        return;
    }

    // Standard API Call
    let apiUrl = 'https://api.cerebras.ai/v1/chat/completions', apiKey = 'csk-mvww3vy29hykeektyv65w9rkjx94hw4r6mrcj5tjcw9942d2';

    console.log(`[DEBUG] Calling API: ${apiUrl} for model: ${apiModelId}`);    const response = await axios.post(apiUrl, { model: apiModelId, messages: apiMessages, stream: true, temperature: 0.7 }, { headers: { 'Authorization': `Bearer ${apiKey}`, 'Content-Type': 'application/json' }, responseType: 'stream' });

    let isStreamEnded = false;

    response.data.on('data', chunk => {
      const lines = chunk.toString().split('\n');
      for (const line of lines) {
        if (line.startsWith('data: ')) {
          const dataStr = line.slice(6).trim();
          if (dataStr === '[DONE]') {
            isStreamEnded = true;
            if (!res.writableEnded) {
                res.write(`data: ${JSON.stringify({ done: true, sessionId: session._id })}\n\n`);
                res.end();
            }
            return;
          }
          try {
            const data = JSON.parse(dataStr);
            const content = data.choices[0].delta?.content || "";
            if (content) {
              fullAIResponse += content;
              res.write(`data: ${JSON.stringify({ message: content })}\n\n`);
            }
          } catch (e) {}
        }
      }
    });

    response.data.on('end', async () => {
      if (res.writableEnded) return;
      await Message.create({ sessionId: session._id, sender: 'user', content: message || "[SIGNAL]", attachmentUrl });
      await Message.create({ sessionId: session._id, sender: 'ai', content: fullAIResponse, modelUsed: activeModelName });
      user.usage.requestsToday += 1;
      await user.save();
      res.write(`data: ${JSON.stringify({ done: true, sessionId: session._id })}\n\n`);
      res.end();
    });

    response.data.on('error', (err) => {
      console.error("Stream Error:", err);
      if (!res.writableEnded) {
        res.write(`data: ${JSON.stringify({ error: "STREAM_ERROR", details: err.message })}\n\n`);
        res.end();
      }
    });

  } catch (err) {
    console.error("[DEBUG] Chat Controller Catch:", err.message);
    if (!res.writableEnded) {
      // Only send error if we haven't sent any AI content yet
      if (!fullAIResponse || fullAIResponse.length < 5) {
        res.write(`data: ${JSON.stringify({ error: "NEURAL_LINK_SEVERED", details: err.message })}\n\n`);
      }
      res.write(`data: ${JSON.stringify({ done: true })}\n\n`);
      res.end();
    }
  }
});

exports.getModels = asyncHandler(async (req, res, next) => { res.status(200).json({ success: true, data: Object.keys(MODELS) }); });
exports.getSessions = asyncHandler(async (req, res, next) => { const sessions = await ChatSession.find({ userId: req.user.id }).sort({ updatedAt: -1 }); res.status(200).json({ success: true, data: sessions }); });
exports.getSessionMessages = asyncHandler(async (req, res, next) => { const messages = await Message.find({ sessionId: req.params.id }).sort({ createdAt: 1 }); res.status(200).json({ success: true, data: messages }); });
exports.clearHistory = asyncHandler(async (req, res, next) => { const sessions = await ChatSession.find({ userId: req.user.id }); const sessionIds = sessions.map(s => s._id); await Message.deleteMany({ sessionId: { $in: sessionIds } }); await ChatSession.deleteMany({ userId: req.user.id }); res.status(200).json({ success: true, data: {} }); });