import express from "express";
import cors from "cors";
import path from "path";
import { fileURLToPath } from "url";
import dotenv from "dotenv";
import { pipeline } from "@huggingface/transformers";

dotenv.config();

const __filename = fileURLToPath(import.meta.url);
const __dirname = path.dirname(__filename);

// --- CUSTOMIZE YOUR AI HERE ---
const SYSTEM_PROMPT = "You are a highly advanced AI assistant named Phi-3 Explorer. Your goal is to provide precise, helpful, and technically accurate responses. Always maintain a professional and sophisticated tone.";
// ------------------------------

let generator: any = null;

async function getGenerator() {
  if (!generator) {
    console.log("🚀 Initializing Phi-3 Core... (Downloading weights ~2.3GB)");
    // Xenova/Phi-3-mini-4k-instruct is optimized for local/edge inference
    generator = await pipeline('text-generation', 'Xenova/Phi-3-mini-4k-instruct', {
      device: 'cpu', 
    });
    console.log("✅ Phi-3 Core Online.");
  }
  return generator;
}

async function startServer() {
  const app = express();
  const PORT = process.env.PORT || 3000; // Use HF provided port or 3000

  app.use(cors());
  app.use(express.json());

  // API Route for Phi-3 Chat
  app.post("/api/chat", async (req, res) => {
    const { messages } = req.body;

    try {
      const gen = await getGenerator();
      
      // Build the prompt with the System Instruction
      let prompt = `<|system|>\n${SYSTEM_PROMPT}<|end|>\n`;
      
      prompt += messages.map((m: any) => {
        const role = m.role === 'user' ? 'user' : 'assistant';
        return `<|${role}|>\n${m.content}<|end|>`;
      }).join("\n") + "\n<|assistant|>";

      console.log("🤖 Generating response...");
      const output = await gen(prompt, {
        max_new_tokens: 1024,
        temperature: 0.7,
        do_sample: true,
        return_full_text: false,
      });

      let text = output[0].generated_text;
      text = text.replace(/<\|end\|>/g, "").trim();

      res.json({ message: { role: "assistant", content: text } });
    } catch (error: any) {
      console.error("❌ Inference Error:", error);
      res.status(500).json({ error: "The AI core encountered an error during generation." });
    }
  });

  app.listen(PORT, "0.0.0.0", async () => {
    console.log(`📡 Backend listening on port ${PORT}`);
    // Start pre-loading the model immediately on boot
    try {
      await getGenerator();
    } catch (e) {
      console.error("Failed to pre-load model:", e);
    }
  });
}

startServer();