import express from 'express';
import { SystemMessage, HumanMessage, Runnable, LlamaCppLLM } from './src/index.js';
import bodyParser from 'body-parser';
import path from 'path';
import fs from 'fs';

// Classify Logic
class EmailClassifierRunnable extends Runnable {
    constructor(llm) {
        super();
        this.llm = llm;
    }
    async _call(input, config) {
        // Mock fallback if model fails
        if (!this.llm) return { category: "Error", confidence: 0, reason: "Model not initialized" };
        
        const messages = this._buildPrompt(input);
        const response = await this.llm.invoke(messages, config);
        return this._parseClassification(response.content);
    }
    _buildPrompt(email) {
        return [
            new SystemMessage(`You are an email classification assistant. Classify into: Spam, Invoice, Meeting Request, Urgent, Personal, Other. Respond in JSON like {"category": "X", "confidence": 0.9, "reason": "Y"}.`),
            new HumanMessage(`Classify:\nSubject: ${email.subject}\nBody: ${email.body}`)
        ];
    }
    _parseClassification(response) {
        try {
            const jsonMatch = response.match(/\{[\s\S]*\}/);
            if (!jsonMatch) throw new Error('No JSON found');
            return JSON.parse(jsonMatch[0]);
        } catch (e) { return { category: 'Other', confidence: 0, reason: 'Failed to parse JSON', raw: response }; }
    }
}

const app = express();
const PORT = 7860;

app.use(bodyParser.json());

let classifier = null;

async function initModel() {
    try {
        console.log("Loading model...");
        // Ensure model exists
        if (!fs.existsSync('./models/Qwen3-1.7B-Q8_0.gguf')) {
            console.error("Model file missing!");
            return;
        }

        const llm = new LlamaCppLLM({
            modelPath: './models/Qwen3-1.7B-Q8_0.gguf',
            temperature: 0.1,
            maxTokens: 200
        });
        
        // Warmup
        await llm.invoke("Hi");
        
        classifier = new EmailClassifierRunnable(llm);
        console.log("Model loaded successfully!");
    } catch (err) {
        console.error("Failed to load model:", err);
    }
}

app.post('/classify', async (req, res) => {
    if (!classifier) return res.status(503).json({ error: "Model loading or failed" });
    try {
        const { subject, body } = req.body;
        const result = await classifier.invoke({ subject, body, from: 'api' });
        res.json(result);
    } catch (error) {
        res.status(500).json({ error: error.message });
    }
});

app.get('/', (req, res) => res.send('AI Email Classifier Running. POST /classify to use.'));

app.listen(PORT, '0.0.0.0', () => {
    console.log(`Server listening on ${PORT}`);
    initModel();
});