const express = require("express");
const http = require("http");
const path = require("path");
const multer = require("multer");
const fs = require("fs").promises;
const { Server } = require("socket.io");
const { ChatGroq } = require("@langchain/groq");
const { HumanMessage, SystemMessage } = require("@langchain/core/messages");
const mammoth = require("mammoth");
const pdf = require("pdf-parse");
const Tesseract = require("tesseract.js");
const sharp = require("sharp");
const cors = require("cors");

const app = express();
const server = http.createServer(app);
const io = new Server(server, {
  cors: { origin: "*" },
  maxHttpBufferSize: 1e8
});

app.use(cors());
app.use(express.json());
app.use(express.static(path.resolve("./public")));
app.use("/uploads", express.static(path.join(__dirname, "uploads")));

// Configure file upload
const storage = multer.diskStorage({
  destination: async (req, file, cb) => {
    const uploadDir = path.join(__dirname, "uploads");
    await fs.mkdir(uploadDir, { recursive: true });
    cb(null, uploadDir);
  },
  filename: (req, file, cb) => {
    const uniqueName = `${Date.now()}-${file.originalname}`;
    cb(null, uniqueName);
  }
});

const upload = multer({ 
  storage,
  limits: { fileSize: 50 * 1024 * 1024 }
});

// Initialize Groq LLM
const llm = new ChatGroq({
  model: "llama-3.3-70b-versatile",
  temperature: 0.7,
  maxTokens: 2000,
  maxRetries: 2,
  apiKey: process.env.GROQ_API_KEY
});


// Data structures
let rooms = {};
let users = {};

// 🚨 FEATURE 4: Emergency keywords detection
const EMERGENCY_KEYWORDS = [
  'chest pain', 'heart attack', 'can\'t breathe', 'breathless', 'severe bleeding',
  'unconscious', 'stroke', 'paralysis', 'severe headache', 'suicide',
  'overdose', 'seizure', 'choking', 'anaphylaxis', 'severe pain'
];

// 🎯 FEATURE 1: Dynamic Dual-Persona AI Safety Engine
const PATIENT_AI_PROMPT = `You are an AI Medical Assistant helping a PATIENT. Your role:

**SAFETY-FIRST APPROACH**
1. **Empathetic Support**: Be warm, reassuring, and supportive
2. **Simple Language**: Avoid medical jargon, explain in simple terms
3. **Symptom Clarification**: Ask ONE focused question at a time
4. **No Premature Conclusions**: Never diagnose or interpret lab results
5. **Safety Boundaries**: If critical values detected, advise immediate medical attention
6. **Respond Only When**:
   - Patient asks direct questions
   - Patient is alone and needs guidance
   - Someone mentions @ai

**RISK CONTROL**: Never share detailed medical analysis. Acknowledge uploads and reassure.`;

const DOCTOR_AI_PROMPT = `You are an AI Medical Assistant helping a DOCTOR. Your role:

**CLINICAL-GRADE ANALYSIS**
1. **Detailed Insights**: Provide comprehensive medical analysis
2. **Critical Findings**: Highlight abnormal values, red flags with clinical context
3. **Medical Terminology**: Use appropriate professional language
4. **Evidence-Based**: Reference standard clinical thresholds
5. **Explainable AI**: Always explain WHY a finding is significant
6. **Respond Only When**:
   - Doctor asks about files/reports
   - Doctor mentions @ai
   - Doctor needs clinical summary

**TRANSPARENCY**: Provide clear reasoning for all flagged findings with confidence levels.`;

// 🔬 FEATURE 3: Explainable AI Layer
async function analyzeFileWithXAI(content, fileName, previousReports = []) {
  const analysisPrompt = `Analyze this medical report with EXPLAINABLE AI principles:

File: ${fileName}
Content: ${content.substring(0, 3000)}

${previousReports.length > 0 ? `
**TEMPORAL CONTEXT** (Previous Reports):
${previousReports.map((r, i) => `Report ${i+1} (${r.date}): ${r.keyFindings}`).join('\n')}
` : ''}

Provide analysis in this EXACT format:

**CLINICAL SUMMARY**
• Main diagnosis/finding (1 line)

**CRITICAL FINDINGS**
• [Value/Finding]: [Normal Range] → [Current Value] → [Deviation %]
  Reason: [Clinical explanation]
  Confidence: [High/Medium/Low]

**TEMPORAL TRENDS** (if previous data available)
• [Parameter]: [Previous → Current] → [Trend Analysis]

**IMMEDIATE CONCERNS**
• [Priority level]: [Specific concern]

**RECOMMENDATIONS**
• [Actionable next steps]

Be concise, clinical, and ALWAYS explain the "why" behind findings.`;

  try {
    const analysis = await llm.invoke([
      new SystemMessage("You are a clinical AI analyzer specializing in explainable medical insights."),
      new HumanMessage(analysisPrompt)
    ]);
    return analysis.content;
  } catch (error) {
    console.error("XAI Analysis error:", error);
    return "Unable to analyze with full explainability.";
  }
}

// 🕐 FEATURE 2: Temporal Health Intelligence
function extractTemporalData(room) {
  if (!room.files || room.files.length < 2) return [];
  
  return room.files.map(f => ({
    name: f.name,
    date: f.uploadedAt,
    keyFindings: f.analysis ? f.analysis.substring(0, 200) : "No analysis",
    content: f.content.substring(0, 500)
  }));
}

async function performTemporalAnalysis(currentContent, fileName, room) {
  const previousReports = extractTemporalData(room);
  
  if (previousReports.length === 0) {
    return await analyzeFileWithXAI(currentContent, fileName, []);
  }

  const temporalPrompt = `Perform TEMPORAL HEALTH INTELLIGENCE analysis:

**CURRENT REPORT**: ${fileName}
${currentContent.substring(0, 2000)}

**HISTORICAL DATA**:
${previousReports.map((r, i) => `
Report ${i+1} - ${new Date(r.date).toLocaleDateString()}:
${r.keyFindings}
`).join('\n')}

Analyze:
1. **Longitudinal Trends**: Compare current vs historical values
2. **Progression/Deterioration**: Identify gradual changes over time
3. **Early Warning Signs**: Flag subtle patterns that indicate future risk
4. **Clinical Significance**: Is this progression normal or concerning?

Format as structured clinical analysis with temporal context.`;

  try {
    const analysis = await llm.invoke([
      new SystemMessage("You are a temporal medical intelligence analyzer specializing in longitudinal health trends."),
      new HumanMessage(temporalPrompt)
    ]);
    return analysis.content;
  } catch (error) {
    console.error("Temporal analysis error:", error);
    return await analyzeFileWithXAI(currentContent, fileName, previousReports);
  }
}

// 🚨 FEATURE 4: Emergency Detection and Escalation
async function detectEmergency(message, userRole) {
  const messageLower = message.toLowerCase();
  
  // Check for emergency keywords
  const hasEmergencyKeyword = EMERGENCY_KEYWORDS.some(keyword => 
    messageLower.includes(keyword)
  );

  if (!hasEmergencyKeyword) return { isEmergency: false };

  // Enhanced AI-based emergency detection
  const emergencyPrompt = `Analyze this message for medical emergency indicators:

Message: "${message}"

Classify emergency level:
- CRITICAL: Immediate life threat (chest pain, can't breathe, severe bleeding, stroke symptoms)
- HIGH: Urgent medical attention needed within hours
- MODERATE: Medical evaluation needed soon
- LOW: Non-emergency concern

Respond ONLY with JSON:
{
  "level": "CRITICAL|HIGH|MODERATE|LOW",
  "reasoning": "brief explanation",
  "urgentAdvice": "immediate action to take"
}`;

  try {
    const response = await llm.invoke([
      new SystemMessage("You are an emergency medical triage AI. Respond ONLY with valid JSON."),
      new HumanMessage(emergencyPrompt)
    ]);

    const result = JSON.parse(response.content.replace(/```json|```/g, '').trim());
    
    return {
      isEmergency: result.level === "CRITICAL" || result.level === "HIGH",
      level: result.level,
      reasoning: result.reasoning,
      urgentAdvice: result.urgentAdvice
    };
  } catch (error) {
    console.error("Emergency detection error:", error);
    return { isEmergency: hasEmergencyKeyword, level: "HIGH", reasoning: "Keyword detected" };
  }
}

// 📋 FEATURE 5: Doctor Co-Pilot Documentation
async function generateClinicalDocumentation(roomId) {
  const room = rooms[roomId];
  if (!room) return null;

  const conversationHistory = room.messages
    .filter(m => m.role === 'Patient' || m.role === 'Doctor')
    .map(m => `${m.role}: ${m.content}`)
    .join('\n');

  const filesSummary = room.files
    .map(f => `- ${f.name}: ${f.analysis || 'No analysis'}`)
    .join('\n');

  const docPrompt = `Generate structured clinical documentation from this consultation:

**CONVERSATION**:
${conversationHistory}

**UPLOADED FILES**:
${filesSummary}

Generate SOAP NOTE format:

**SUBJECTIVE**
- Chief Complaint: [main issue]
- History of Present Illness: [brief narrative]
- Review of Systems: [relevant findings]

**OBJECTIVE**
- Vital signs/Reports: [from uploaded files]
- Physical findings: [mentioned in chat]

**ASSESSMENT**
- Primary diagnosis: [clinical impression]
- Differential diagnoses: [alternatives]

**PLAN**
- Investigations: [tests ordered]
- Treatment: [medications/interventions]
- Follow-up: [next steps]

Keep concise and clinically accurate.`;

  try {
    const documentation = await llm.invoke([
      new SystemMessage("You are a medical documentation AI specializing in SOAP notes and clinical summaries."),
      new HumanMessage(docPrompt)
    ]);
    return documentation.content;
  } catch (error) {
    console.error("Documentation generation error:", error);
    return null;
  }
}

// Helper: OCR for images
async function extractTextFromImage(imagePath) {
  try {
    console.log("Starting OCR:", imagePath);
    const processedPath = imagePath + "_processed.jpg";
    await sharp(imagePath)
      .greyscale()
      .normalize()
      .sharpen()
      .toFile(processedPath);

    const { data: { text } } = await Tesseract.recognize(processedPath, 'eng');
    
    try { await fs.unlink(processedPath); } catch (e) {}
    
    console.log("OCR completed, text length:", text.length);
    return text.trim();
  } catch (error) {
    console.error("OCR Error:", error);
    return "";
  }
}

// Helper: Extract text from files
async function extractFileContent(filePath, mimeType) {
  try {
    console.log("Extracting:", filePath, mimeType);
    
    if (mimeType === "application/pdf") {
      const dataBuffer = await fs.readFile(filePath);
      const pdfData = await pdf(dataBuffer);
      return pdfData.text;
    } else if (mimeType.includes("word") || mimeType.includes("document")) {
      const result = await mammoth.extractRawText({ path: filePath });
      return result.value;
    } else if (mimeType.includes("text")) {
      return await fs.readFile(filePath, "utf-8");
    } else if (mimeType.includes("image")) {
      const ocrText = await extractTextFromImage(filePath);
      return ocrText.length > 10 ? ocrText : "[Image - no text detected]";
    }
    return "[Unsupported format]";
  } catch (error) {
    console.error("Extraction error:", error);
    return "[Extraction failed]";
  }
}

// Helper: AI Response with risk-aware disclosure control
async function getAIResponse(roomId, userMessage, userRole, isFileQuery = false, emergencyContext = null) {
  const room = rooms[roomId];
  if (!room) return "Room not found";

  // FEATURE 1: Dynamic persona selection
  const systemPrompt = userRole === "doctor" ? DOCTOR_AI_PROMPT : PATIENT_AI_PROMPT;
  
  const roleMessages = room.messages.filter(m => 
    !m.forRole || m.forRole === userRole || (!m.forRole && m.role !== 'AI Assistant')
  );
  
  let context = `Room: ${roomId}
User Role: ${userRole}
Patient: ${room.patient || "Waiting"}
Doctor: ${room.doctor || "Not yet joined"}

${emergencyContext ? `🚨 EMERGENCY CONTEXT: ${emergencyContext.reasoning}\nLevel: ${emergencyContext.level}` : ''}

Recent messages (last 5):
${roleMessages.slice(-5).map(m => `${m.role}: ${m.content}`).join("\n")}`;

  // FEATURE 1: Risk-based information disclosure
  if (userRole === "doctor" && isFileQuery && room.files.length > 0) {
    context += `\n\n**CLINICAL FILES** (with XAI explanations):\n${room.files.map((f, i) => 
      `${i+1}. ${f.name}\n   Analysis: ${f.analysis}\n   Key content: ${f.content.substring(0, 400)}`
    ).join("\n\n")}`;
  } else if (userRole === "patient" && room.files.length > 0) {
    // Patients get minimal, safe information
    context += `\n\n**FILES UPLOADED**: ${room.files.map(f => f.name).join(', ')}
Note: Detailed medical analysis is being reviewed by your doctor.`;
  }

  const messages = [
    new SystemMessage(systemPrompt),
    new SystemMessage(context),
    new HumanMessage(`[${userRole}]: ${userMessage}`)
  ];

  try {
    const response = await llm.invoke(messages);
    return response.content;
  } catch (error) {
    console.error("AI Error:", error);
    return "I'm having trouble responding. Please try again.";
  }
}

// File upload endpoint with TEMPORAL ANALYSIS
app.post("/upload", upload.single("file"), async (req, res) => {
  try {
    const { roomId, uploadedBy, uploaderRole } = req.body;
    const file = req.file;

    if (!file || !roomId) {
      return res.status(400).json({ error: "File and roomId required" });
    }

    console.log("Upload:", file.originalname, "by", uploadedBy, "in", roomId);

    const content = await extractFileContent(file.path, file.mimetype);
    console.log("Content extracted, length:", content.length);

    // FEATURE 2 & 3: Temporal analysis with XAI
    let analysis = "";
    if (content && content.length > 20 && !content.includes("no text detected")) {
      if (rooms[roomId]) {
        analysis = await performTemporalAnalysis(content, file.originalname, rooms[roomId]);
      } else {
        analysis = await analyzeFileWithXAI(content, file.originalname, []);
      }
    }

    const fileInfo = {
      name: file.originalname,
      path: file.path,
      url: `/uploads/${file.filename}`,
      type: file.mimetype,
      content: content.substring(0, 5000),
      analysis: analysis,
      uploadedAt: new Date().toISOString(),
      uploadedBy: uploadedBy || "Unknown"
    };

    if (rooms[roomId]) {
      rooms[roomId].files.push(fileInfo);
      
      // Broadcast file upload to everyone in room
      const fileMessage = {
        role: uploadedBy || "User",
        nickname: uploadedBy,
        content: `📎 Uploaded: ${file.originalname}`,
        timestamp: new Date().toISOString(),
        fileData: {
          name: file.originalname,
          url: fileInfo.url,
          type: file.mimetype,
          analysis: analysis
        },
        isFile: true
      };

      rooms[roomId].messages.push(fileMessage);
      io.to(roomId).emit("chat-message", fileMessage);
      
      // Emit file list update to all users in the room
      io.to(roomId).emit("files-updated", { files: rooms[roomId].files });

      // FEATURE 1: Role-specific AI responses (PRIVATE - not visible to other role)
      if (content && content.length > 20) {
        setTimeout(() => {
          const doctorSocketId = Object.keys(users).find(
            sid => users[sid].roomId === roomId && users[sid].role === "doctor"
          );
          
          if (doctorSocketId && rooms[roomId].doctor) {
            const doctorAiMessage = `🔬 **Clinical Analysis** (with XAI)\n\n${analysis}`;
            io.to(doctorSocketId).emit("ai-message", { 
              message: doctorAiMessage, 
              isPrivate: true,
              forRole: "doctor"
            });
          }
        }, 1000);
      }

      if (uploaderRole === "patient") {
        setTimeout(() => {
          const patientSocketId = Object.keys(users).find(
            sid => users[sid].nickname === uploadedBy && users[sid].roomId === roomId
          );
          
          if (patientSocketId) {
            const patientAiMessage = `✅ I've received "${file.originalname}". Your doctor will review it shortly.`;
            io.to(patientSocketId).emit("ai-message", { 
              message: patientAiMessage, 
              isPrivate: true,
              forRole: "patient"
            });
          }
        }, 500);
      }
    }

    res.json({ success: true, file: fileInfo });
  } catch (error) {
    console.error("Upload error:", error);
    res.status(500).json({ error: "Upload failed: " + error.message });
  }
});

// FEATURE 5: Generate clinical documentation endpoint
app.post("/generate-documentation", async (req, res) => {
  try {
    const { roomId } = req.body;
    if (!roomId || !rooms[roomId]) {
      return res.status(400).json({ error: "Invalid room ID" });
    }

    const documentation = await generateClinicalDocumentation(roomId);
    res.json({ success: true, documentation });
  } catch (error) {
    console.error("Documentation error:", error);
    res.status(500).json({ error: "Documentation generation failed" });
  }
});

// Socket.IO
io.on("connection", (socket) => {
  console.log("Connected:", socket.id);

  socket.on("join-room", async ({ roomId, nickname, role }) => {
    socket.join(roomId);
    users[socket.id] = { nickname, role, roomId };

    if (!rooms[roomId]) {
      rooms[roomId] = {
        patient: null,
        doctor: null,
        messages: [],
        files: [],
        patientData: {},
        emergencyMode: false
      };
    }

    if (role === "patient" && !rooms[roomId].patient) {
      rooms[roomId].patient = nickname;
    } else if (role === "doctor" && !rooms[roomId].doctor) {
      rooms[roomId].doctor = nickname;
    }

    socket.emit("room-history", {
      messages: rooms[roomId].messages.filter(m => !m.forRole),
      files: rooms[roomId].files
    });

    io.to(roomId).emit("user-joined", { 
      nickname, 
      role,
      patient: rooms[roomId].patient,
      doctor: rooms[roomId].doctor
    });

    // Role-specific greeting (PRIVATE - only to this user)
    let greeting = "";
    if (role === "patient") {
      greeting = `Hello ${nickname}! 👋 I'm here to help guide you. What brings you in today?`;
    } else if (role === "doctor") {
      greeting = `Welcome Dr. ${nickname}! 👨‍⚕️ Clinical analysis tools ready. Use "Generate SOAP Note" for documentation.`;
      
      // FEATURE 5: Doctor briefing
      if (rooms[roomId].messages.length > 0 || rooms[roomId].files.length > 0) {
        setTimeout(async () => {
          const briefing = await getAIResponse(
            roomId,
            "Provide a 3-point clinical summary: chief complaint, temporal trends from files, critical findings.",
            "doctor",
            true
          );
          
          socket.emit("ai-message", { 
            message: `📋 **Clinical Briefing**:\n${briefing}`, 
            isPrivate: true,
            forRole: "doctor"
          });
        }, 1000);
      }
    }

    if (greeting) {
      socket.emit("ai-message", { 
        message: greeting, 
        isPrivate: true,
        forRole: role
      });
    }
  });

  socket.on("chat-message", async ({ roomId, message }) => {
    const user = users[socket.id];
    if (!user || !rooms[roomId]) return;

    // Check if this is an @ai request
    const isAIRequest = message.toLowerCase().includes('@ai');

    // FEATURE 4: Emergency detection
    const emergencyCheck = await detectEmergency(message, user.role);

    // If NOT an @ai request, broadcast message to everyone
    if (!isAIRequest) {
      const chatMessage = {
        role: user.role === "patient" ? "Patient" : "Doctor",
        nickname: user.nickname,
        content: message,
        timestamp: new Date().toISOString(),
        isEmergency: emergencyCheck.isEmergency
      };

      rooms[roomId].messages.push(chatMessage);
      io.to(roomId).emit("chat-message", chatMessage);
    }

    // FEATURE 4: Emergency escalation
    if (emergencyCheck.isEmergency) {
      rooms[roomId].emergencyMode = true;
      
      // Alert patient immediately
      if (user.role === "patient") {
        const urgentMessage = `🚨 **URGENT MEDICAL ATTENTION NEEDED**\n\n${emergencyCheck.urgentAdvice}\n\nCall emergency services (911) immediately if symptoms worsen.`;
        socket.emit("ai-message", { 
          message: urgentMessage, 
          isPrivate: true,
          forRole: "patient",
          isEmergency: true
        });
      }

      // Alert doctor
      const doctorSocketId = Object.keys(users).find(
        sid => users[sid].roomId === roomId && users[sid].role === "doctor"
      );
      
      if (doctorSocketId) {
        const doctorAlert = `🚨 **EMERGENCY ALERT**\n\nPatient: ${user.nickname}\nLevel: ${emergencyCheck.level}\nReason: ${emergencyCheck.reasoning}\n\nMessage: "${message}"\n\nImmediate evaluation required.`;
        io.to(doctorSocketId).emit("ai-message", { 
          message: doctorAlert, 
          isPrivate: true,
          forRole: "doctor",
          isEmergency: true
        });
      }

      return; // Don't process normal AI response in emergency
    }

    // Handle @ai requests - PRIVATE response only to requester
    if (isAIRequest) {
      const messageText = message.toLowerCase();
      const isFileQuery = 
        messageText.includes("report") || 
        messageText.includes("file") ||
        messageText.includes("result") ||
        messageText.includes("test") ||
        messageText.includes("value") ||
        messageText.includes("finding") ||
        messageText.includes("trend");

      setTimeout(async () => {
        const aiResponse = await getAIResponse(roomId, message, user.role, isFileQuery);
        
        // Send ONLY to the user who requested (not broadcast)
        socket.emit("ai-message", { 
          message: aiResponse, 
          isPrivate: true,
          forRole: user.role
        });
      }, 1500);
    } else {
      // Auto-respond logic for non-@ai messages
      const messageText = message.toLowerCase();
      const isFileQuery = 
        messageText.includes("report") || 
        messageText.includes("file") ||
        messageText.includes("result") ||
        messageText.includes("test") ||
        messageText.includes("value") ||
        messageText.includes("finding") ||
        messageText.includes("trend");

      const shouldAIRespond = 
        (user.role === "patient" && !rooms[roomId].doctor && message.endsWith("?")) ||
        (user.role === "doctor" && isFileQuery);

      if (shouldAIRespond) {
        setTimeout(async () => {
          const aiResponse = await getAIResponse(roomId, message, user.role, isFileQuery);
          
          socket.emit("ai-message", { 
            message: aiResponse, 
            isPrivate: true,
            forRole: user.role
          });
        }, 1500);
      }
    }
  });

  // FEATURE 5: Generate documentation on request
  socket.on("request-documentation", async ({ roomId }) => {
    const user = users[socket.id];
    if (!user || user.role !== "doctor") return;

    const documentation = await generateClinicalDocumentation(roomId);
    if (documentation) {
      socket.emit("documentation-generated", { documentation });
    }
  });

  socket.on("typing", ({ roomId }) => {
    const user = users[socket.id];
    if (user) {
      socket.to(roomId).emit("user-typing", { nickname: user.nickname });
    }
  });

  socket.on("disconnect", () => {
    const user = users[socket.id];
    if (user) {
      const { roomId, nickname, role } = user;
      
      if (rooms[roomId]) {
        if (role === "patient") rooms[roomId].patient = null;
        if (role === "doctor") rooms[roomId].doctor = null;

        io.to(roomId).emit("user-left", { 
          nickname, 
          role,
          patient: rooms[roomId].patient,
          doctor: rooms[roomId].doctor
        });
      }

      delete users[socket.id];
    }
  });
});

const PORT = process.env.PORT || 7860;
server.listen(PORT, "0.0.0.0", () =>
  console.log(`🏥 Enhanced Medical Chat Server running on port ${PORT}`)
);