package org.texttechnologylab.project.gruppe_05_1.nlp;

import com.mongodb.client.model.Filters;
import com.mongodb.client.model.UpdateOneModel;
import com.mongodb.client.model.WriteModel;
import org.apache.uima.fit.util.JCasUtil;
import org.bson.Document;
import java.io.*;
import java.nio.charset.StandardCharsets;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.zip.*;
import java.util.*;
import java.util.stream.Collectors;
import java.util.zip.GZIPInputStream;
import org.bson.conversions.Bson;
import org.apache.uima.fit.factory.JCasFactory;
import org.apache.uima.jcas.JCas;
import org.apache.uima.cas.impl.XmiCasDeserializer;
import org.texttechnologylab.project.gruppe_05_1.database.MongoDBHandler;
import org.texttechnologylab.project.gruppe_05_1.util.Logger;

import static org.texttechnologylab.project.gruppe_05_1.Main.JCAS_SPEECHES_TYPESYSTEM_DIR;


public class XmiExtractor {

    private final List<WriteModel<Document>> bulkOperations = Collections.synchronizedList(new ArrayList<>());
    private final MongoDBHandler mongoDBHandler;
    private static final int BATCH_SIZE = 1000;
    private static final AtomicInteger processedCount = new AtomicInteger(0);

    public XmiExtractor() {
        mongoDBHandler = new MongoDBHandler();
    }

    public void extractAndUploadXmiData() throws IOException {
        InputStream resourceStream = getClass().getClassLoader().getResourceAsStream("speeches/20.zip");
        if (resourceStream == null) {
            throw new IOException("20.zip nicht gefunden im Ressourcenordner /speeches");
        }
        ExecutorService executor = Executors.newFixedThreadPool(Runtime.getRuntime().availableProcessors());
        List<Future<?>> futures = new ArrayList<>();
        try (ZipInputStream zis = new ZipInputStream(resourceStream)) {
            ZipEntry entry;
            while ((entry = zis.getNextEntry()) != null) {
                if (entry.getName().endsWith(".xmi.gz")) {
                    File tempFile = File.createTempFile("xmi_entry_", ".xmi.gz");
                    try (FileOutputStream fos = new FileOutputStream(tempFile)) {
                        byte[] buffer = new byte[1024];
                        int len;
                        while ((len = zis.read(buffer)) > 0) {
                            fos.write(buffer, 0, len);
                        }
                    }
                    ZipEntry finalEntry = entry;
                    Future<?> future = executor.submit(() -> {
                        try (FileInputStream fis = new FileInputStream(tempFile)) {
                            processXmiGzStream(fis, finalEntry.getName());
                        } catch (IOException e) {
                            e.printStackTrace();
                        } finally {
                            tempFile.delete();
                        }
                    });
                    futures.add(future);
                }
                zis.closeEntry();
            }
        } catch (IOException | NullPointerException e) {
            Logger.error("Error reading XMI data from ZIP file.");
            Logger.error(e.getMessage());
        }
        for (Future<?> future : futures) {
            try {
                future.get();
            } catch (Exception e) {
                e.printStackTrace();
            }
        }
        executor.shutdown();
        flushBatch(); // Synchronously upload the remaining batch
        mongoDBHandler.close();
    }

    private void processXmiGzStream(InputStream inputStream, String filename) {
        JCas jCas = null;
        try (GZIPInputStream gis = new GZIPInputStream(inputStream)) {
            jCas = JCasFactory.createJCas(JCAS_SPEECHES_TYPESYSTEM_DIR);
            XmiCasDeserializer.deserialize(gis, jCas.getCas(), true);

            Document analysisResults = new Document();

            List<Document> tokens = new ArrayList<>();
            for (de.tudarmstadt.ukp.dkpro.core.api.segmentation.type.Token token : JCasUtil.select(jCas, de.tudarmstadt.ukp.dkpro.core.api.segmentation.type.Token.class)) {
                Document tokenDoc = new Document()
                        .append("text", token.getCoveredText())
                        .append("pos", token.getPos().getPosValue())
                        .append("lemma", token.getLemma().getValue());
                tokens.add(tokenDoc);
            }
            analysisResults.append("tokens", tokens);

            List<String> sentences = JCasUtil.select(jCas, de.tudarmstadt.ukp.dkpro.core.api.segmentation.type.Sentence.class).stream()
                    .map(de.tudarmstadt.ukp.dkpro.core.api.segmentation.type.Sentence::getCoveredText)
                    .collect(Collectors.toList());
            analysisResults.append("sentences", sentences);

            List<Document> dependencies = new ArrayList<>();
            for (de.tudarmstadt.ukp.dkpro.core.api.syntax.type.dependency.Dependency dep : JCasUtil.select(jCas, de.tudarmstadt.ukp.dkpro.core.api.syntax.type.dependency.Dependency.class)) {
                Document depDoc = new Document()
                        .append("type", dep.getDependencyType())
                        .append("governor", dep.getGovernor().getCoveredText())
                        .append("dependent", dep.getDependent().getCoveredText());
                dependencies.add(depDoc);
            }
            analysisResults.append("dependencies", dependencies);

            List<Document> namedEntities = new ArrayList<>();
            for (de.tudarmstadt.ukp.dkpro.core.api.ner.type.NamedEntity ne : JCasUtil.select(jCas, de.tudarmstadt.ukp.dkpro.core.api.ner.type.NamedEntity.class)) {
                Document neDoc = new Document()
                        .append("text", ne.getCoveredText())
                        .append("type", ne.getValue());
                namedEntities.add(neDoc);
            }
            analysisResults.append("namedEntities", namedEntities);

            List<Document> sentiments = new ArrayList<>();
            for (org.hucompute.textimager.uima.type.Sentiment sentiment : JCasUtil.select(jCas, org.hucompute.textimager.uima.type.Sentiment.class)) {
                Document sentimentDoc = new Document()
                        .append("begin", sentiment.getBegin())
                        .append("end", sentiment.getEnd())
                        .append("score", sentiment.getSentiment())
                        .append("subjectivity", sentiment.getSubjectivity());
                if (sentiment instanceof org.hucompute.textimager.uima.type.VaderSentiment) {
                    org.hucompute.textimager.uima.type.VaderSentiment vader = (org.hucompute.textimager.uima.type.VaderSentiment) sentiment;
                    sentimentDoc.append("pos", vader.getPos())
                            .append("neu", vader.getNeu())
                            .append("neg", vader.getNeg());
                }
                sentiments.add(sentimentDoc);
            }
            analysisResults.append("sentiments", sentiments);

            List<Document> topics = new ArrayList<>();
            for (org.hucompute.textimager.uima.type.category.CategoryCoveredTagged topic : JCasUtil.select(jCas, org.hucompute.textimager.uima.type.category.CategoryCoveredTagged.class)) {
                Document topicDoc = new Document()
                        .append("topic", topic.getValue())
                        .append("score", topic.getScore())
                        .append("text", topic.getCoveredText());
                topics.add(topicDoc);
            }
            topics.sort((d1, d2) -> Double.compare(d2.getDouble("score"), d1.getDouble("score")));
            analysisResults.append("topics", topics);

            String speechKey = extractSpeechKeyFromFilename(filename);
            if (speechKey != null) {
                Bson filter = Filters.eq("speechKey", speechKey);
                Bson update = new Document("$set", new Document("analysisResults", analysisResults));
                UpdateOneModel<Document> updateModel = new UpdateOneModel<>(filter, update);
                bulkOperations.add(updateModel);
                synchronized (bulkOperations) {
                    if (bulkOperations.size() >= BATCH_SIZE) {
                        Logger.info("BATCH_SIZE to Upload: " + bulkOperations.size());
                        flushBatch();
                    }
                }
                int count = processedCount.incrementAndGet();
                if (count % 1000 == 0) {
                    Logger.info("Processed speeches: " + count);
                }
            }
        } catch (Exception e) {
            e.printStackTrace();
        } finally {
            if (jCas != null) {
                jCas.reset();
            }
        }
    }

    private static String extractSpeechKeyFromFilename(String filename) {
        String baseName = filename.replace(".xmi.gz", "");
        return baseName.replace("20/", "");
    }

    private synchronized void flushBatch() {
        if (!bulkOperations.isEmpty()) {
            mongoDBHandler.bulkWriteNlpData(bulkOperations);
            bulkOperations.clear();
        }
    }
}

