package ru.nbch.credit_tracker.service.xml.stax;

import org.apache.commons.lang3.exception.ExceptionUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import ru.nbch.credit_tracker.exceptions.SignalClientException;

import javax.xml.stream.XMLInputFactory;
import javax.xml.stream.XMLStreamException;
import javax.xml.stream.XMLStreamReader;
import java.io.ByteArrayInputStream;
import java.io.InputStream;
import java.io.StringReader;
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.function.Consumer;
import java.util.function.Function;
import java.util.function.Supplier;

public class XmlProcessorStaxImpl implements StaxXmlProcessor {
    private final Logger log = LoggerFactory.getLogger(getClass());
    protected XMLInputFactory xmlInputFactory;
    protected Map<Class<?>, Supplier<StaxReader<?>>> staxReaders = new HashMap<>();

    public XmlProcessorStaxImpl(List<Supplier<StaxReader<?>>> staxReaders) {
        this.xmlInputFactory = XMLInputFactory.newInstance();
        this.xmlInputFactory.setProperty(XMLInputFactory.SUPPORT_DTD, false); // полностью выключить DOCTYPE Declaration
        // или выключить особо опасные возможности (менее строгое ограничение XXE, но достаточное):
        this.xmlInputFactory.setProperty(XMLInputFactory.IS_SUPPORTING_EXTERNAL_ENTITIES, false);
        this.xmlInputFactory.setProperty(XMLInputFactory.IS_NAMESPACE_AWARE, false);
        this.xmlInputFactory.setXMLResolver((publicID, systemID, baseURI, namespace) -> {
            throw new XMLStreamException("External entities are disabled");
        });
        staxReaders.forEach(supplier -> {
            Class<?> type = supplier.get().type();
            this.staxReaders.put(type, supplier);
        });
    }

    public XmlProcessorStaxImpl(Supplier<StaxReader<?>>... staxReaders) {
        this(Arrays.asList(staxReaders));
    }

    @Override
    public <Z> Z read(String xml, Class<Z> clazz) {
        try {
            XMLStreamReader xmlStreamReader = xmlInputFactory.createXMLStreamReader(new StringReader(xml));
            Supplier<StaxReader<?>> staxReaderSupplier = staxReaders.get(clazz);
            @SuppressWarnings("unchecked")
            StaxReader<Z> staxProcessor = (StaxReader<Z>) staxReaderSupplier.get();
            if (staxProcessor != null) {
                return staxProcessor.read(xmlStreamReader);
            }
            throw new IllegalStateException("no stax reader registered for " + clazz);
        } catch (XMLStreamException e) {
            log.warn(ExceptionUtils.getStackTrace(e));
            throw new RuntimeException(e);
        } catch (Exception e) {
            log.error(ExceptionUtils.getStackTrace(e));
            throw new RuntimeException(e);
        }
    }

    @Override
    public <Z> Z read(byte[] xml, Class<Z> clazz) {
        try {
            XMLStreamReader xmlStreamReader = xmlInputFactory.createXMLStreamReader(new ByteArrayInputStream(xml));
            Supplier<StaxReader<?>> staxReaderSupplier = staxReaders.get(clazz);
            @SuppressWarnings("unchecked")
            StaxReader<Z> staxProcessor = (StaxReader<Z>) staxReaderSupplier.get();
            if (staxProcessor != null) {
                return staxProcessor.read(xmlStreamReader);
            }
            throw new IllegalStateException("no stax reader registered for " + clazz);
        } catch (XMLStreamException e) {
            log.warn(ExceptionUtils.getStackTrace(e));
            throw new RuntimeException(e);
        } catch (Exception e) {
            log.error(ExceptionUtils.getStackTrace(e));
            throw new RuntimeException(e);
        }
    }

    @Override
    public <Z> Z read(InputStream inputStream, Class<Z> clazz) {
        try {
            XMLStreamReader xmlStreamReader = xmlInputFactory.createXMLStreamReader(inputStream);
            Supplier<StaxReader<?>> staxReaderSupplier = staxReaders.get(clazz);
            @SuppressWarnings("unchecked")
            StaxReader<Z> staxProcessor = (StaxReader<Z>) staxReaderSupplier.get();
            if (staxProcessor != null) {
                return staxProcessor.read(xmlStreamReader);
            }
            throw new IllegalStateException("no stax reader registered for " + clazz);
        } catch (XMLStreamException e) {
            log.warn(ExceptionUtils.getStackTrace(e));
            throw new RuntimeException(e);
        } catch (Exception e) {
            log.error(ExceptionUtils.getStackTrace(e));
            throw new RuntimeException(e);
        }
    }

    @Override
    public <Z> Z read(InputStream inputStream, Class<Z> clazz, Consumer<Z> objProcessor, Function<Z, Boolean> emptyBatchCheckFunction, Integer batchSize, String endBatchElementPath) {
        try {
            XMLStreamReader xmlStreamReader = xmlInputFactory.createXMLStreamReader(inputStream);
            Supplier<StaxReader<?>> staxReaderSupplier = staxReaders.get(clazz);
            @SuppressWarnings("unchecked")
            StaxReader<Z> staxProcessor = (StaxReader<Z>) staxReaderSupplier.get();
            if (staxProcessor != null) {
                return staxProcessor.read(xmlStreamReader, objProcessor, emptyBatchCheckFunction, batchSize, endBatchElementPath);
            }
            throw new IllegalStateException("no stax reader registered for " + clazz);
        } catch (SignalClientException e) {
            throw e;
        } catch (XMLStreamException e) {
            log.warn(ExceptionUtils.getStackTrace(e));
            throw new RuntimeException(e);
        } catch (Exception e) {
            log.error(ExceptionUtils.getStackTrace(e));
            throw new RuntimeException(e);
        }
    }
}
