package org.texttechnologylab.project.gruppe_05_1.nlp;

import com.mongodb.client.MongoCollection;
import com.mongodb.client.MongoDatabase;
import com.mongodb.client.model.Filters;
import com.mongodb.client.model.UpdateOneModel;
import com.mongodb.client.model.WriteModel;
import com.mongodb.client.result.UpdateResult;
import org.apache.uima.fit.util.JCasUtil;
import org.bson.Document;
import java.io.*;
import java.net.URL;
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 de.tudarmstadt.ukp.dkpro.core.api.segmentation.type.Token;
import de.tudarmstadt.ukp.dkpro.core.api.segmentation.type.Sentence;
import de.tudarmstadt.ukp.dkpro.core.api.syntax.type.dependency.Dependency;
import de.tudarmstadt.ukp.dkpro.core.api.ner.type.NamedEntity;
import org.texttechnologylab.project.gruppe_05_1.database.MongoDBHandler;
import org.hucompute.textimager.uima.type.category.CategoryCoveredTagged;
import org.texttechnologylab.project.gruppe_05_1.util.Logger;


public class XmiExtractor {

    private List<WriteModel<Document>> bulkOperations;
    private MongoDBHandler mongoDBHandler;
    private static final int BATCH_SIZE = 1000;
    private int processedCount = 0;
    public XmiExtractor() {
        mongoDBHandler = new MongoDBHandler();
        this.bulkOperations = new ArrayList<>();
    }

    public void extractAndUploadXmiData() throws IOException {
        InputStream resourceStream = getClass().getClassLoader().getResourceAsStream("speeches/20.zip");
        try (ZipInputStream zis = new ZipInputStream(resourceStream)) {
            ZipEntry entry;
            while ((entry = zis.getNextEntry()) != null) {
                if (entry.getName().endsWith(".xmi.gz")) {
                    ByteArrayOutputStream baos = new ByteArrayOutputStream();
                    byte[] buffer = new byte[1024];
                    int len;
                    while ((len = zis.read(buffer)) > 0) {
                        baos.write(buffer, 0, len);
                    }
                    byte[] entryData = baos.toByteArray();
                    processXmiGzStream(new ByteArrayInputStream(entryData), entry.getName());
                }
                zis.closeEntry();
            }
        } catch (IOException | NullPointerException e) {
            Logger.error("Error reading XMI data from ZIP file.");
            Logger.error(e.getMessage());
        }
        flushBatch();
    }

    private void processXmiGzStream(InputStream inputStream, String filename) {
        JCas jCas = null;
        try (GZIPInputStream gis = new GZIPInputStream(inputStream)) {
            jCas = JCasFactory.createJCas("src/main/resources/speeches/TypeSystem");
            XmiCasDeserializer.deserialize(gis, jCas.getCas(), true);

            // Build structured analysisResults Document
            Document analysisResults = new Document();

            // Tokens: Include POS, Lemma, etc.
            List<Document> tokens = new ArrayList<>();
            for (Token token : JCasUtil.select(jCas, 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);

            // Sentences
            List<String> sentences = JCasUtil.select(jCas, Sentence.class).stream()
                    .map(Sentence::getCoveredText)
                    .collect(Collectors.toList());
            analysisResults.append("sentences", sentences);

            // Dependencies
            List<Document> dependencies = new ArrayList<>();
            for (Dependency dep : JCasUtil.select(jCas, 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);

            // Named Entities
            List<Document> namedEntities = new ArrayList<>();
            for (NamedEntity ne : JCasUtil.select(jCas, NamedEntity.class)) {
                Document neDoc = new Document()
                        .append("text", ne.getCoveredText())
                        .append("type", ne.getValue());
                namedEntities.add(neDoc);
            }
            analysisResults.append("namedEntities", namedEntities);

            // Sentiment
            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());

                // Check for VaderSentiment subtype
                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 (CategoryCoveredTagged topic : JCasUtil.select(jCas, CategoryCoveredTagged.class)) {
                Document topicDoc = new Document()
                        .append("topic", topic.getValue())
                        .append("score", topic.getScore())
                        .append("tags", topic.getTags())
                        .append("text", topic.getCoveredText());
                topics.add(topicDoc);
            }
            topics.sort((d1, d2) -> Double.compare(d2.getDouble("score"), d1.getDouble("score")));
            analysisResults.append("topics", topics);


            // Upload structured Document to MongoDB
            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);
                if (bulkOperations.size() >= BATCH_SIZE) {
                    flushBatch();
                }
                processedCount++;
                if (processedCount % 5000 == 0) {
                    System.out.println("Processed speeches: " + processedCount);
                }

            }

        } 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 void flushBatch() {
        if (!bulkOperations.isEmpty()) {
            mongoDBHandler.bulkWriteNlpData(bulkOperations);
            bulkOperations.clear();
        }
    }


    /*
    public static void main(String[] args) {
        try {
            XmiExtractor extractor = new XmiExtractor(database);
            extractor.extractAndUploadXmiData();
            System.out.println("Processing complete.");
        } catch (Exception e) {
            e.printStackTrace();
        }
    } */
}

