import express from 'express';
import session from 'express-session';
import passport from 'passport';
import { Strategy as GoogleStrategy } from 'passport-google-oauth20';
import pgSession from 'connect-pg-simple';
import pool from './config/db.js';
import dotenv from 'dotenv';
import cors from 'cors';
import multer from 'multer';
import path from 'path';
import fs from 'fs';
import { Ollama } from 'ollama';
import { createRequire } from 'module';
const customRequire = createRequire(import.meta.url);
const parsePDF = customRequire('pdf-parse');
const { readFile, utils } = customRequire('xlsx');
const AdmZip = customRequire('adm-zip');
import mammoth from 'mammoth';
import { UserProfile } from './types/index.js';
import { fileURLToPath } from 'url';

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

dotenv.config();

const app = express();
const port = process.env.PORT || 3000;
const PostgresStore = pgSession(session);
const ollama = new Ollama({ host: process.env.OLLAMA_HOST });

// Load User Profiles
const userProfilesPath = path.join(process.cwd(), 'user_profiles.json');
const userProfiles = JSON.parse(fs.readFileSync(userProfilesPath, 'utf-8')).user_profiles;

// Multer Setup
const upload = multer({
    storage: multer.diskStorage({
        destination: (req, file, cb) => {
            const uploadDir = process.env.UPLOAD_DIR || '/mnt/data';
            if (!fs.existsSync(uploadDir)) {
                fs.mkdirSync(uploadDir, { recursive: true });
            }
            cb(null, uploadDir);
        },
        filename: (req, file, cb) => {
            const uniqueSuffix = Date.now() + '-' + Math.round(Math.random() * 1e9);
            cb(null, `${uniqueSuffix}-${file.originalname}`);
        },
    }),
});

// Middleware
const corsOptions = process.env.NODE_ENV === 'production'
    ? { origin: 'https://ai.dffm.it', credentials: true }
    : { origin: true, credentials: true };
app.use(cors(corsOptions));
app.use(express.json());
app.use(session({
    store: new PostgresStore({
        pool: pool,
        tableName: 'session'
    }),
    secret: process.env.SESSION_SECRET || 'secret',
    resave: false,
    saveUninitialized: false,
    cookie: { maxAge: 30 * 24 * 60 * 60 * 1000 } // 30 days
}));

app.use(passport.initialize());
app.use(passport.session());

// Serve static files from React app
app.use(express.static(path.join(__dirname, '../../client/dist')));

// Passport Config
passport.use(new GoogleStrategy({
    clientID: process.env.GOOGLE_CLIENT_ID!,
    clientSecret: process.env.GOOGLE_CLIENT_SECRET!,
    callbackURL: process.env.CALLBACK_URL!,
},
    async (accessToken, refreshToken, profile, done) => {
        const email = profile.emails?.[0].value;
        if (!email) return done(new Error('No email found'));

        const userProfile = userProfiles[email];

        if (userProfile) {
            return done(null, { ...userProfile, email });
        } else {
            // Deny access or assign guest as requested
            // I'll deny for security unless user is in json
            return done(null, false, { message: 'Unauthorized' });
        }
    }
));

passport.serializeUser((user: any, done) => {
    done(null, user);
});

passport.deserializeUser((user: any, done) => {
    done(null, user);
});

// Auth Routes
app.get('/auth/google', passport.authenticate('google', { scope: ['profile', 'email'] }));

app.get('/auth/google/callback',
    passport.authenticate('google', { failureRedirect: '/login' }),
    (req, res) => {
        res.redirect('/api/me');
    }
);

app.post('/auth/logout', (req, res, next) => {
    req.logout((err) => {
        if (err) { return next(err); }
        req.session.destroy((err) => {
            if (err) {
                console.error('Session destruction error:', err);
            }
            res.clearCookie('connect.sid');
            res.json({ message: 'Logged out successfully' });
        });
    });
});

// Helper to ensure authenticated
const isAuthenticated = (req: express.Request, res: express.Response, next: express.NextFunction) => {
    if (req.isAuthenticated()) return next();
    res.status(401).json({ error: 'Unauthorized' });
};

// API Endpoints
app.get('/api/me', isAuthenticated, (req, res) => {
    res.json(req.user);
});

app.get('/api/documents', isAuthenticated, async (req, res) => {
    try {
        const user = req.user as UserProfile;
        const result = await pool.query(
            "SELECT DISTINCT filename, MAX(id) as id FROM documents WHERE uploaded_by = $1 GROUP BY filename ORDER BY id DESC",
            [user.email]
        );
        res.json(result.rows);
    } catch (error: any) {
        console.error(error);
        res.status(500).json({ error: error.message });
    }
});

// Helper to chunk text
const chunkText = (text: string, size: number = 1000, overlap: number = 200): string[] => {
    const chunks: string[] = [];
    for (let i = 0; i < text.length; i += size - overlap) {
        chunks.push(text.slice(i, i + size));
    }
    return chunks;
};

// Task 4: Ingest
app.post('/api/ingest', isAuthenticated, upload.single('file'), async (req, res) => {
    try {
        if (!req.file) return res.status(400).json({ error: 'No file uploaded' });

        const user = req.user as UserProfile;
        const filePath = req.file.path;
        // Pre-cleaning: remove old versions of this file for the same user/collection
        await pool.query(
            "DELETE FROM documents WHERE filename = $1 AND uploaded_by = $2",
            [req.file.originalname, user.email]
        );

        const ext = path.extname(req.file.originalname).toLowerCase();
        let fullContent = "";

        // Extraction Strategies
        if (ext === '.pdf') {
            const dataBuffer = fs.readFileSync(filePath);
            const data = await parsePDF(dataBuffer);
            fullContent = data.text;
        } else if (ext === '.docx') {
            const data = await mammoth.extractRawText({ path: filePath });
            fullContent = data.value;
        } else if (ext === '.odt') {
            try {
                const zip = new AdmZip(filePath);
                const contentXml = zip.readAsText("content.xml");
                if (!contentXml) throw new Error("Could not read content.xml from ODT");

                // Better ODT XML parsing: find all text:p and text:h tags
                const textMatch = contentXml.match(/<text:[ph][^>]*>([\s\S]*?)<\/text:[ph]>/g);
                if (textMatch) {
                    fullContent = textMatch.map((m: string) => m.replace(/<[^>]+>/g, '')).join("\n");
                } else {
                    // Fallback to simple tag stripping if complex match fails
                    fullContent = contentXml.replace(/<[^>]+>/g, ' ');
                }
            } catch (err) {
                console.error("ODT Error:", err);
                throw new Error("Failed to parse ODT file correctly");
            }
        } else if (['.xlsx', '.xls', '.csv', '.ods'].includes(ext)) {
            try {
                const workbook = readFile(filePath);
                const sheetName = workbook.SheetNames[0];
                const sheet = workbook.Sheets[sheetName];
                const rows = utils.sheet_to_json(sheet);

                // Process each row as a separate "semantic unit"
                for (let i = 0; i < rows.length; i++) {
                    const row = rows[i] as any;
                    const machineId = row.TCO || row.Matricola || row.SerialNumber || `Riga ${i + 1}`;
                    const machineType = row["Machine Type"] || row.Machine || row.Risorsa || "N/A";

                    const rowText = Object.entries(row)
                        .map(([k, v]) => `${k}: ${v}`)
                        .join(" | ");

                    // Strong context header for each row
                    const contextualizedRow = `[DOC: ${req.file.originalname}] [MACCHINA: ${machineId}] [TIPO: ${machineType}] DATI: ${rowText}`;

                    const sanitizedRow = contextualizedRow.replace(/\0/g, '').replace(/[^\x00-\x7F]/g, ' ').replace(/\s+/g, ' ').trim();
                    if (!sanitizedRow) continue;

                    const embeddingResponse = await ollama.embeddings({
                        model: 'nomic-embed-text',
                        prompt: sanitizedRow,
                    });

                    await pool.query(
                        `INSERT INTO documents (path, filename, content, embedding, collection, uploaded_by)
                         VALUES ($1, $2, $3, $4, $5, $6)`,
                        [`${filePath}#row${i}`, req.file.originalname, sanitizedRow, JSON.stringify(embeddingResponse.embedding), user.rag_collection, user.email]
                    );
                }
                res.json({ message: `Excel/ODS ingested successfully: ${rows.length} rows as separate units`, path: filePath });
                return; // Prevent fallthrough to general chunking logic
            } catch (err) {
                console.error("XLSX Error:", err);
                throw new Error("Failed to parse Spreadsheet file correctly");
            }
        } else {
            // Default as text (txt, md, etc)
            fullContent = fs.readFileSync(filePath, 'utf-8');
        }

        // Clean content for token efficiency and remove NULL/Non-UTF8-friendly bytes for Postgres
        fullContent = fullContent.replace(/\0/g, '').replace(/[^\x00-\x7F]/g, ' ').replace(/\s+/g, ' ').trim();

        // Split into chunks to avoid context limits
        const chunks = chunkText(fullContent, 1000, 200);
        console.log(`Ingesting ${req.file.originalname} (${ext}): ${chunks.length} chunks`);

        for (let i = 0; i < chunks.length; i++) {
            const chunk = chunks[i];
            const embeddingResponse = await ollama.embeddings({
                model: 'nomic-embed-text',
                prompt: chunk,
            });

            const embedding = embeddingResponse.embedding;

            await pool.query(
                `INSERT INTO documents (path, filename, content, embedding, collection, uploaded_by)
                 VALUES ($1, $2, $3, $4, $5, $6)`,
                [`${filePath}#chunk${i}`, req.file.originalname, chunk, JSON.stringify(embedding), user.rag_collection, user.email]
            );
        }

        res.json({ message: `Document ingested successfully: ${chunks.length} chunks`, path: filePath });
    } catch (error: any) {
        console.error(error);
        res.status(500).json({ error: error.message });
    }
});

// Task 4: Search
app.post('/api/search', isAuthenticated, async (req, res) => {
    try {
        const { query } = req.body;
        if (!query) return res.status(400).json({ error: 'Query is required' });

        const user = req.user as UserProfile;

        // Generate Query Embedding
        const embeddingResponse = await ollama.embeddings({
            model: 'nomic-embed-text',
            prompt: query,
        });
        const embedding = embeddingResponse.embedding;

        // Hybrid Search: Similarity score + Keyword Boost (multiplier)
        const result = await pool.query(
            `SELECT id, filename, content, 
             (1 - (embedding <=> $1::vector)) * (CASE WHEN content ILIKE $2 THEN 2.0 ELSE 1.0 END) as similarity
             FROM documents
             WHERE collection = $3
             AND (content ILIKE $2 OR embedding <=> $1::vector < 0.6)
             ORDER BY similarity DESC
             LIMIT 10`,
            [JSON.stringify(embedding), `%${query}%`, user.rag_collection]
        );

        res.json(result.rows);
    } catch (error: any) {
        console.error(error);
        res.status(500).json({ error: error.message });
    }
});

// Task 2: Implement /api/chat (RAG Endpoint)
// Database Schema Migration
const initDb = async () => {
    const client = await pool.connect();
    try {
        await client.query('BEGIN');

        // Workspaces
        await client.query(`
            CREATE TABLE IF NOT EXISTS workspaces (
                id SERIAL PRIMARY KEY,
                user_email VARCHAR(255) NOT NULL,
                name VARCHAR(255) NOT NULL,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            );
        `);

        // Chats
        await client.query(`
            CREATE TABLE IF NOT EXISTS chats (
                id SERIAL PRIMARY KEY,
                workspace_id INTEGER REFERENCES workspaces(id),
                user_email VARCHAR(255) NOT NULL,
                title VARCHAR(255),
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            );
        `);

        // Messages
        await client.query(`
            CREATE TABLE IF NOT EXISTS messages (
                id SERIAL PRIMARY KEY,
                chat_id INTEGER REFERENCES chats(id),
                role VARCHAR(50) NOT NULL, -- 'user' or 'assistant'
                content TEXT NOT NULL,
                sources TEXT[],
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            );
        `);

        // Ensure pgvector extension
        await client.query('CREATE EXTENSION IF NOT EXISTS vector;');

        // Documents (Existing, ensuring structure)
        await client.query(`
            CREATE TABLE IF NOT EXISTS documents (
                id SERIAL PRIMARY KEY,
                user_email VARCHAR(255),
                filename VARCHAR(255),
                content TEXT,
                embedding vector(768),
                collection VARCHAR(50),
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            );
        `);

        await client.query('COMMIT');
        console.log('Database initialized successfully');
    } catch (e) {
        await client.query('ROLLBACK');
        console.error('Database initialization failed:', e);
    } finally {
        client.release();
    }
};

initDb();

// ... (Existing Middleware and Routes)

// Task 1: Fetch Ollama Models
app.get('/api/models', isAuthenticated, async (req, res) => {
    try {
        const response = await fetch(`${process.env.OLLAMA_HOST || 'http://localhost:11434'}/api/tags`);
        if (!response.ok) throw new Error('Failed to fetch models from Ollama');
        const data = await response.json();
        const models = data.models.map((m: any) => m.name);
        res.json(models);
    } catch (error: any) {
        console.error('Model fetch error:', error);
        res.status(500).json({ error: 'Failed to fetch models' });
    }
});

// New Endpoint: Get Workspaces/Folders
app.get('/api/workspaces', isAuthenticated, async (req, res) => {
    const user = req.user as UserProfile;
    try {
        const result = await pool.query('SELECT * FROM workspaces WHERE user_email = $1 ORDER BY created_at DESC', [user.email]);
        // Return default if empty
        if (result.rows.length === 0) {
            res.json([{ id: 0, name: 'General', user_email: user.email }]);
        } else {
            res.json(result.rows);
        }
    } catch (err: any) {
        res.status(500).json({ error: err.message });
    }
});

// New Endpoint: Get Chats
app.get('/api/chats', isAuthenticated, async (req, res) => {
    const user = req.user as UserProfile;
    try {
        const result = await pool.query('SELECT * FROM chats WHERE user_email = $1 ORDER BY created_at DESC LIMIT 50', [user.email]);
        res.json(result.rows);
    } catch (err: any) {
        res.status(500).json({ error: err.message });
    }
});

// New Endpoint: Get Messages for a Chat
app.get('/api/chats/:chatId/messages', isAuthenticated, async (req, res) => {
    try {
        const { chatId } = req.params;
        const result = await pool.query('SELECT * FROM messages WHERE chat_id = $1 ORDER BY created_at ASC', [chatId]);
        res.json(result.rows);
    } catch (err: any) {
        res.status(500).json({ error: err.message });
    }
});

// Task 2: Implement /api/chat (RAG Endpoint with Toggle & Persistence)
app.post('/api/chat', isAuthenticated, async (req, res) => {
    try {
        const { message, model = 'llama3:8b', useRag = true, chatId: requestedChatId } = req.body;
        if (!message) return res.status(400).json({ error: 'Message is required' });

        const user = req.user as UserProfile;

        // 0. Ensure Chat Exists
        let chatId = requestedChatId;
        if (!chatId) {
            const chatResult = await pool.query(
                'INSERT INTO chats (user_email, title) VALUES ($1, $2) RETURNING id',
                [user.email, message.substring(0, 50)]
            );
            chatId = chatResult.rows[0].id;
        }

        // Save User Message
        await pool.query(
            'INSERT INTO messages (chat_id, role, content) VALUES ($1, $2, $3)',
            [chatId, 'user', message]
        );

        let prompt = message;
        let sources: string[] = [];

        if (useRag) {
            // 1. Embed Query
            const embeddingResponse = await ollama.embeddings({
                model: 'nomic-embed-text',
                prompt: message,
            });
            const embedding = embeddingResponse.embedding;

            // Hybrid Search: Similarity score + Multi-term Keyword Boost
            const terms: string[] = message.split(/\s+/).filter((t: string) => t.length >= 2);
            const keywordQuery = terms.length > 0
                ? terms.map((t: string) => `content ILIKE '%${t.replace(/'/g, "''")}%'`).join(' AND ')
                : '1=1';

            const dbResult = await pool.query(
                `SELECT filename, content, 
                 (1 - (embedding <=> $1::vector)) * (CASE WHEN ${keywordQuery} THEN 3.0 ELSE 1.0 END) as similarity
                 FROM documents
                 WHERE collection = $2
                 ORDER BY similarity DESC
                 LIMIT 15`,
                [JSON.stringify(embedding), user.rag_collection]
            );

            const context = dbResult.rows.map(row => row.content).join("\n\n---\n\n");
            sources = Array.from(new Set(dbResult.rows.map(row => row.filename)));

            // 3. Generazione risposta via llama3:8b (o modello selezionato)
            prompt = `
            Sei un ESTRATTORE DATI TECNICI per sistemi industriali.
            Il tuo obiettivo è fornire dati precisi e distinguere chiaramente i diversi macchinari.

            REGOLE DI RISPOSTA (MANDATORIE):
            1. LINGUA: Rispondi esclusivamente in ITALIANO.
            2. AMBIGUITÀ: Se nel CONTESTO trovi più record che corrispondono alla domanda (es. diversi IP per macchine diverse), ELENCALI TUTTI specificando la Macchina/Matricola di appartenenza. NON scegliere un valore a caso.
            3. FORMATO: Usa elenchi puntati o tabelle. Sii schematico e tecnico.
            4. NO CHIAACCHIERE: Non iniziare con "Ecco i dati" o "Spero di aiutare". Vai dritto al dato.
            5. ZERO INVENZIONE: Se non trovi il dato esatto per la macchina richiesta, dì "Dato non trovato per [Nome Macchina]".

            CONTESTO DISPONIBILE:
            ${context}

            DOMANDA UTENTE:
            ${message}

            RISPOSTA TECNICA (Italiano):
            `;
        }

        const response = await ollama.generate({
            model: model,
            prompt: prompt,
            stream: false,
        });

        const answer = response.response;

        // Save Assistant Message
        await pool.query(
            'INSERT INTO messages (chat_id, role, content, sources) VALUES ($1, $2, $3, $4)',
            [chatId, 'assistant', answer, sources]
        );

        res.json({
            chatId,
            answer,
            sources
        });

    } catch (error: any) {
        console.error(error);
        res.status(500).json({ error: error.message });
    }
});

// SPA Fallback: serve index.html for any other requests
app.get(/.*/, (req, res) => {
    res.sendFile(path.join(__dirname, '../../client/dist/index.html'));
});

app.listen(port, () => {
    console.log(`Server running at http://192.168.1.239:${port}`);
});
