package io.xpipe.app.util;

import com.fasterxml.jackson.annotation.JsonIdentityInfo;
import com.fasterxml.jackson.annotation.JsonPropertyOrder;
import com.fasterxml.jackson.annotation.ObjectIdGenerators;
import com.fasterxml.jackson.databind.annotation.JsonSerialize;
import io.xpipe.app.beacon.BeaconAuthMethod;
import io.xpipe.app.beacon.BeaconClientInformation;
import io.xpipe.app.cred.SecurityKeyImpl;
import io.xpipe.app.cred.SshIdentityStrategy;
import io.xpipe.app.ext.HostAddress;
import io.xpipe.app.process.ShellDialect;
import io.xpipe.app.process.ShellDialects;
import io.xpipe.app.process.ShellScript;
import io.xpipe.app.pwman.*;
import io.xpipe.app.rdp.ExternalRdpClient;
import io.xpipe.app.secret.*;
import io.xpipe.app.spice.ExternalSpiceClient;
import io.xpipe.app.storage.*;
import io.xpipe.app.terminal.ExternalTerminalType;
import io.xpipe.app.terminal.TerminalMultiplexer;
import io.xpipe.app.terminal.TerminalPrompt;
import io.xpipe.app.vnc.ExternalVncClient;

import com.fasterxml.jackson.core.JsonGenerator;
import com.fasterxml.jackson.core.JsonParser;
import com.fasterxml.jackson.databind.*;
import com.fasterxml.jackson.databind.deser.ContextualDeserializer;
import com.fasterxml.jackson.databind.jsontype.NamedType;
import com.fasterxml.jackson.databind.jsontype.TypeDeserializer;
import com.fasterxml.jackson.databind.jsontype.TypeSerializer;
import com.fasterxml.jackson.databind.jsontype.impl.AsPropertyTypeDeserializer;
import com.fasterxml.jackson.databind.module.SimpleModule;
import com.fasterxml.jackson.databind.node.JsonNodeFactory;
import com.fasterxml.jackson.databind.node.ObjectNode;
import com.fasterxml.jackson.databind.type.SimpleType;

import java.io.CharArrayReader;
import java.io.IOException;
import java.nio.charset.Charset;
import java.nio.file.InvalidPathException;
import java.nio.file.Path;
import java.util.ArrayList;
import java.util.List;
import java.util.UUID;
import java.util.stream.Stream;

public class AppJacksonModule extends SimpleModule {

    @Override
    public void setupModule(SetupContext context) {
        context.registerSubtypes(
                new NamedType(BeaconClientInformation.Api.class),
                new NamedType(BeaconClientInformation.Cli.class),
                new NamedType(BeaconClientInformation.Daemon.class));
        context.registerSubtypes(
                new NamedType(BeaconAuthMethod.Local.class), new NamedType(BeaconAuthMethod.ApiKey.class));

        context.registerSubtypes(VaultKeySecretValue.class);
        context.registerSubtypes(PasswordLockSecretValue.class);

        addSerializer(DataStoreEntryRef.class, new DataStoreEntryRefSerializer());
        addDeserializer(DataStoreEntryRef.class, new DataStoreEntryRefDeserializer());
        addSerializer(ContextualFileReference.class, new LocalFileReferenceSerializer());
        addDeserializer(ContextualFileReference.class, new LocalFileReferenceDeserializer());
        addSerializer(ExternalTerminalType.class, new ExternalTerminalTypeSerializer());
        addDeserializer(ExternalTerminalType.class, new ExternalTerminalTypeDeserializer());
        addSerializer(EncryptedValue.class, new EncryptedValueSerializer());
        addDeserializer(EncryptedValue.class, new EncryptedValueDeserializer<>());
        addSerializer(EncryptedValue.CurrentKey.class, new EncryptedValueSerializer());
        addDeserializer(EncryptedValue.CurrentKey.class, new EncryptedValueDeserializer<>());
        addSerializer(EncryptedValue.VaultKey.class, new EncryptedValueSerializer());
        addDeserializer(EncryptedValue.VaultKey.class, new EncryptedValueDeserializer<>());

        addSerializer(ShellDialect.class, new ShellDialectSerializer());
        addDeserializer(ShellDialect.class, new ShellDialectDeserializer());

        addSerializer(OsType.class, new OsTypeSerializer());
        addDeserializer(OsType.Local.class, new OsTypeLocalDeserializer());
        addDeserializer(OsType.Any.class, new OsTypeAnyDeserializer());

        addSerializer(ShellScript.class, new ShellScriptSerializer());
        addDeserializer(ShellScript.class, new ShellScriptDeserializer());

        addSerializer(HostAddress.class, new HostAddressSerializer());
        addDeserializer(HostAddress.class, new HostAddressDeserializer());

        addSerializer(KeePassXcPasswordManager.class, new KeePassXcPasswordManagerSerializer());
        addDeserializer(KeePassXcPasswordManager.class, new KeePassXcPasswordManagerDeserializer());

        for (ShellDialect t : ShellDialects.ALL) {
            context.registerSubtypes(new NamedType(t.getClass()));
        }

        context.registerSubtypes(SecurityKeyImpl.getClasses());
        context.registerSubtypes(SshIdentityStrategy.getClasses());
        context.registerSubtypes(PasswordManagerKeyStrategy.getClasses());
        context.registerSubtypes(PasswordManager.getClasses());
        context.registerSubtypes(TerminalMultiplexer.getClasses());
        context.registerSubtypes(TerminalPrompt.getClasses());
        context.registerSubtypes(ExternalVncClient.getClasses());
        context.registerSubtypes(ExternalRdpClient.getClasses());
        context.registerSubtypes(ExternalSpiceClient.getClasses());
        context.registerSubtypes(SecretRetrievalStrategy.getClasses());
        context.registerSubtypes(DataStorageGroupStrategy.getClasses());
        context.registerSubtypes(KeeperPasswordManager.KeeperAuth.getClasses());

        context.registerSubtypes(new NamedType(InPlaceSecretValue.class));

        addSerializer(FilePath.class, new FilePathSerializer());
        addDeserializer(FilePath.class, new FilePathDeserializer());

        addSerializer(StorePath.class, new StorePathSerializer());
        addDeserializer(StorePath.class, new StorePathDeserializer());

        addSerializer(Charset.class, new CharsetSerializer());
        addDeserializer(Charset.class, new CharsetDeserializer());

        addSerializer(Path.class, new LocalPathSerializer());
        addDeserializer(Path.class, new LocalPathDeserializer());

        context.setMixInAnnotations(Throwable.class, ThrowableTypeMixIn.class);

        super.setupModule(context);
    }

    public static class StorePathSerializer extends JsonSerializer<StorePath> {

        @Override
        public void serialize(StorePath value, JsonGenerator jgen, SerializerProvider provider) throws IOException {
            var ar = value.getNames().toArray(String[]::new);
            jgen.writeArray(ar, 0, ar.length);
        }
    }

    public static class StorePathDeserializer extends JsonDeserializer<StorePath> {

        @Override
        public StorePath deserialize(JsonParser p, DeserializationContext ctxt) throws IOException {
            JavaType javaType =
                    JacksonMapper.getDefault().getTypeFactory().constructCollectionLikeType(List.class, String.class);
            List<String> list = JacksonMapper.getDefault().readValue(p, javaType);
            return new StorePath(list);
        }
    }

    public static class FilePathSerializer extends JsonSerializer<FilePath> {

        @Override
        public void serialize(FilePath value, JsonGenerator jgen, SerializerProvider provider) throws IOException {
            jgen.writeString(value.toString());
        }
    }

    public static class FilePathDeserializer extends JsonDeserializer<FilePath> {

        @Override
        public FilePath deserialize(JsonParser p, DeserializationContext ctxt) throws IOException {
            return FilePath.of(p.getValueAsString());
        }
    }

    public static class CharsetSerializer extends JsonSerializer<Charset> {

        @Override
        public void serialize(Charset value, JsonGenerator jgen, SerializerProvider provider) throws IOException {
            jgen.writeString(value.name());
        }
    }

    public static class CharsetDeserializer extends JsonDeserializer<Charset> {

        @Override
        public Charset deserialize(JsonParser p, DeserializationContext ctxt) throws IOException {
            return Charset.forName(p.getValueAsString());
        }
    }

    public static class LocalPathSerializer extends JsonSerializer<Path> {

        @Override
        public void serialize(Path value, JsonGenerator jgen, SerializerProvider provider) throws IOException {
            jgen.writeString(value.toString());
        }
    }

    public static class LocalPathDeserializer extends JsonDeserializer<Path> {

        @Override
        public Path deserialize(JsonParser p, DeserializationContext ctxt) throws IOException {
            try {
                return Path.of(p.getValueAsString());
            } catch (InvalidPathException ignored) {
                return null;
            }
        }
    }

    @JsonSerialize(as = Throwable.class)
    @JsonPropertyOrder(alphabetic = true)
    public abstract static class ThrowableTypeMixIn {

        @SuppressWarnings("unused")
        @JsonIdentityInfo(generator = ObjectIdGenerators.StringIdGenerator.class, property = "$id")
        private Throwable cause;
    }

    public static class KeePassXcPasswordManagerSerializer extends JsonSerializer<KeePassXcPasswordManager> {

        @Override
        public void serialize(KeePassXcPasswordManager value, JsonGenerator jgen, SerializerProvider provider)
                throws IOException {
            if (value == null) {
                jgen.writeNull();
                return;
            }

            var tree = JacksonMapper.getDefault().valueToTree(value.getAssociationKeys());
            var object = JsonNodeFactory.instance.objectNode();
            object.put("type", "keePassXc");
            object.set("associationKeys", tree);
            object.set("keyStrategy", JacksonMapper.getDefault().valueToTree(value.getKeyStrategy()));
            jgen.writeTree(object);
        }

        @Override
        public void serializeWithType(
                KeePassXcPasswordManager value,
                JsonGenerator gen,
                SerializerProvider serializers,
                TypeSerializer typeSer)
                throws IOException {
            serialize(value, gen, serializers);
        }
    }

    public static class KeePassXcPasswordManagerDeserializer extends JsonDeserializer<KeePassXcPasswordManager> {

        @Override
        @SuppressWarnings("unchecked")
        public KeePassXcPasswordManager deserialize(JsonParser p, DeserializationContext ctxt) throws IOException {
            JsonNode tree = JacksonMapper.getDefault().readTree(p);
            if (tree == null || !tree.isObject()) {
                return null;
            }

            var keyStrategyNode = tree.get("keyStrategy");
            var keyStrategy = keyStrategyNode != null
                    ? JacksonMapper.getDefault().treeToValue(keyStrategyNode, PasswordManagerKeyStrategy.class)
                    : null;

            if (tree.has("associationKey")) {
                var parsed = JacksonMapper.getDefault()
                        .treeToValue(tree.required("associationKey"), KeePassXcAssociationKey.class);
                return KeePassXcPasswordManager.builder()
                        .keyStrategy(keyStrategy)
                        .associationKeys(parsed != null ? List.of(parsed) : List.of())
                        .build();
            } else {
                var javaType = JacksonMapper.getDefault()
                        .getTypeFactory()
                        .constructCollectionLikeType(List.class, KeePassXcAssociationKey.class);
                var parsed = (List<KeePassXcAssociationKey>)
                        JacksonMapper.getDefault().treeToValue(tree.required("associationKeys"), javaType);
                return KeePassXcPasswordManager.builder()
                        .keyStrategy(keyStrategy)
                        .associationKeys(parsed)
                        .build();
            }
        }
    }

    public static class OsTypeSerializer extends JsonSerializer<OsType> {

        @Override
        public void serialize(OsType value, JsonGenerator jgen, SerializerProvider provider) throws IOException {
            jgen.writeString(value.getId());
        }
    }

    public static class OsTypeLocalDeserializer extends JsonDeserializer<OsType.Local> {

        @Override
        public OsType.Local deserialize(JsonParser p, DeserializationContext ctxt) throws IOException {
            var stream = Stream.of(OsType.WINDOWS, OsType.LINUX, OsType.MACOS);
            var n = p.getValueAsString();
            return stream.filter(osType ->
                            osType.getName().equals(n) || osType.getId().equals(n))
                    .findFirst()
                    .orElse(null);
        }
    }

    public static class OsTypeAnyDeserializer extends JsonDeserializer<OsType.Any> {

        @Override
        public OsType.Any deserialize(JsonParser p, DeserializationContext ctxt) throws IOException {
            var stream = Stream.of(
                    OsType.WINDOWS, OsType.LINUX, OsType.BSD, OsType.SOLARIS, OsType.MACOS, OsType.AIX, OsType.UNIX);
            var n = p.getValueAsString();
            return stream.filter(osType ->
                            osType.getName().equals(n) || osType.getId().equals(n))
                    .findFirst()
                    .orElse(null);
        }
    }

    public static class LocalFileReferenceSerializer extends JsonSerializer<ContextualFileReference> {

        @Override
        public void serialize(ContextualFileReference value, JsonGenerator jgen, SerializerProvider provider)
                throws IOException {
            jgen.writeString(value.serialize());
        }
    }

    public static class ShellDialectSerializer extends JsonSerializer<ShellDialect> {

        @Override
        public void serialize(ShellDialect value, JsonGenerator jgen, SerializerProvider provider) throws IOException {
            jgen.writeString(value.getId());
        }
    }

    public static class ShellDialectDeserializer extends JsonDeserializer<ShellDialect> {

        @Override
        public ShellDialect deserialize(JsonParser p, DeserializationContext ctxt) throws IOException {
            JsonNode tree = JacksonMapper.getDefault().readTree(p);
            if (tree.isObject()) {
                var t = tree.get("type");
                if (t == null) {
                    return null;
                }
                return ShellDialects.byIdIfPresent(t.asText()).orElse(null);
            }

            return ShellDialects.byIdIfPresent(tree.asText()).orElse(null);
        }
    }

    public static class ShellScriptSerializer extends JsonSerializer<ShellScript> {

        @Override
        public void serialize(ShellScript value, JsonGenerator jgen, SerializerProvider provider) throws IOException {
            jgen.writeString(value.getValue());
        }
    }

    public static class ShellScriptDeserializer extends JsonDeserializer<ShellScript> {

        @Override
        public ShellScript deserialize(JsonParser p, DeserializationContext ctxt) throws IOException {
            return new ShellScript(p.getValueAsString());
        }
    }

    public static class LocalFileReferenceDeserializer extends JsonDeserializer<ContextualFileReference> {

        @Override
        public ContextualFileReference deserialize(JsonParser p, DeserializationContext ctxt) throws IOException {
            return ContextualFileReference.of(p.getValueAsString());
        }
    }

    public static class ExternalTerminalTypeSerializer extends JsonSerializer<ExternalTerminalType> {

        @Override
        public void serialize(ExternalTerminalType value, JsonGenerator jgen, SerializerProvider provider)
                throws IOException {
            jgen.writeString(value.getId());
        }
    }

    public static class ExternalTerminalTypeDeserializer extends JsonDeserializer<ExternalTerminalType> {

        @Override
        public ExternalTerminalType deserialize(JsonParser p, DeserializationContext ctxt) throws IOException {
            var id = p.getValueAsString();
            return ExternalTerminalType.ALL_ON_ALL_PLATFORMS.stream()
                    .filter(terminalType -> terminalType.getId().equals(id))
                    .findFirst()
                    .orElse(null);
        }
    }

    @SuppressWarnings("all")
    public static class EncryptedValueSerializer extends JsonSerializer<EncryptedValue> {

        @Override
        public void serialize(EncryptedValue value, JsonGenerator jgen, SerializerProvider provider)
                throws IOException {
            if (value.getValue() == null) {
                jgen.writeNull();
                return;
            }

            jgen.writeTree(value.getSecret().serialize(value.allowUserSecretKey()));
        }

        @Override
        public void serializeWithType(
                EncryptedValue value, JsonGenerator gen, SerializerProvider serializers, TypeSerializer typeSer)
                throws IOException {
            if (value.getValue() == null) {
                gen.writeNull();
                return;
            }

            gen.writeTree(value.getSecret().serialize(value.allowUserSecretKey()));
        }
    }

    @SuppressWarnings("all")
    public static class EncryptedValueDeserializer<T extends EncryptedValue<?>> extends JsonDeserializer<T>
            implements ContextualDeserializer {

        private boolean useCurrentSecretKeyIfPossible;
        private boolean forceCurrentSecretKey;
        private Class<?> type;

        @Override
        @SuppressWarnings("unchecked")
        public JsonDeserializer<?> createContextual(DeserializationContext ctxt, BeanProperty property)
                throws JsonMappingException {
            var deserializer = new EncryptedValueDeserializer();
            if (property == null) {
                return deserializer;
            }

            JavaType wrapperType = property.getType();
            JavaType valueType = wrapperType.containedType(0);
            deserializer.useCurrentSecretKeyIfPossible =
                    !wrapperType.getRawClass().equals(EncryptedValue.VaultKey.class);
            deserializer.forceCurrentSecretKey = wrapperType.getRawClass().equals(EncryptedValue.CurrentKey.class);
            deserializer.type = valueType.getRawClass();
            return deserializer;
        }

        @Override
        @SuppressWarnings("unchecked")
        public T deserialize(JsonParser p, DeserializationContext ctxt) throws IOException {
            if (type == null) {
                return null;
            }

            return (T) get(p, type, useCurrentSecretKeyIfPossible, forceCurrentSecretKey);
        }

        public Object deserializeWithType(JsonParser jp, DeserializationContext ctxt, TypeDeserializer typeDeserializer)
                throws IOException {
            var type = ((AsPropertyTypeDeserializer) typeDeserializer).baseType();
            JavaType wrapperType = type;
            JavaType valueType = wrapperType.containedType(0);
            var useCurrentSecretKey = !wrapperType.equals(SimpleType.constructUnsafe(EncryptedValue.VaultKey.class));
            var forceCurrentSecretKey = wrapperType.equals(SimpleType.constructUnsafe(EncryptedValue.CurrentKey.class));
            return get(jp, valueType.getRawClass(), useCurrentSecretKey, forceCurrentSecretKey);
        }

        private EncryptedValue get(
                JsonParser p, Class<?> type, boolean useCurrentSecretKey, boolean forceCurrentSecretKey)
                throws IOException {
            if (forceCurrentSecretKey && DataStorageUserHandler.getInstance().getActiveUser() == null) {
                return null;
            }

            Object value;
            JsonNode tree = JacksonMapper.getDefault().readTree(p);
            var secret = DataStorageSecret.deserialize(tree);
            if (secret == null) {
                var raw = JacksonMapper.getDefault().treeToValue(tree, type);
                if (raw != null) {
                    value = raw;
                    var s = JacksonMapper.getDefault().writeValueAsString(value);
                    var internalSecret = InPlaceSecretValue.of(s.toCharArray());
                    secret = DataStorageSecret.ofSecret(
                            internalSecret,
                            useCurrentSecretKey
                                            && DataStorageUserHandler.getInstance()
                                                            .getActiveUser()
                                                    != null
                                    ? EncryptionToken.ofUser()
                                    : EncryptionToken.ofVaultKey());
                } else {
                    return null;
                }
            } else {
                if (!secret.getEncryptedToken().canDecrypt()) {
                    return null;
                }

                var s = secret.getSecret();
                if (s.length == 0) {
                    return null;
                }
                value = JacksonMapper.getDefault().readValue(new CharArrayReader(s), type);
                if (value == null) {
                    return null;
                }
            }
            var perUser = useCurrentSecretKey;
            return perUser
                    ? new EncryptedValue.CurrentKey<>(value, secret)
                    : new EncryptedValue.VaultKey<>(value, secret);
        }
    }

    @SuppressWarnings("all")
    public static class DataStoreEntryRefSerializer extends JsonSerializer<DataStoreEntryRef> {

        @Override
        public void serialize(DataStoreEntryRef value, JsonGenerator jgen, SerializerProvider provider)
                throws IOException {
            if (value == null) {
                jgen.writeNull();
                return;
            }

            jgen.writeString(value.get().getUuid().toString());
        }
    }

    public static class DataStoreEntryRefDeserializer extends JsonDeserializer<DataStoreEntryRef<?>> {

        @Override
        public DataStoreEntryRef<?> deserialize(JsonParser p, DeserializationContext ctxt) throws IOException {
            JsonNode tree = p.getCodec().readTree(p);
            if (tree == null) {
                return null;
            }

            String text;
            if (tree.isObject()) {
                var obj = (ObjectNode) tree;
                if (!obj.has("storeId") || !obj.required("storeId").isTextual()) {
                    return null;
                }

                text = obj.required("storeId").asText();
                if (text.isBlank()) {
                    return null;
                }
            } else {
                if (!tree.isTextual()) {
                    return null;
                }
                text = tree.asText();
            }

            var id = UUID.fromString(text);
            // Keep an invalid entry if it is per-user, meaning that it will get removed later on
            var e = DataStorage.get()
                    .getStoreEntryIfPresent(id)
                    .filter(dataStoreEntry -> dataStoreEntry.getValidity() != DataStoreEntry.Validity.LOAD_FAILED
                            || !dataStoreEntry.getStoreNode().isReadableForUser())
                    .orElse(null);
            if (e == null) {
                return null;
            }

            return new DataStoreEntryRef<>(e);
        }
    }

    public static class HostAddressSerializer extends JsonSerializer<HostAddress> {

        @Override
        public void serialize(HostAddress value, JsonGenerator jgen, SerializerProvider provider) throws IOException {
            if (value.isSingle()) {
                jgen.writeString(value.get());
            } else {
                var tree = JsonNodeFactory.instance.objectNode();
                tree.put("value", value.get());
                tree.set("available", JacksonMapper.getDefault().valueToTree(value.getAvailable()));
                jgen.writeTree(tree);
            }
        }
    }

    public static class HostAddressDeserializer extends JsonDeserializer<HostAddress> {

        @Override
        public HostAddress deserialize(JsonParser p, DeserializationContext ctxt) throws IOException {
            var tree = (JsonNode) p.getCodec().readTree(p);
            if (tree.isTextual()) {
                return !tree.textValue().isBlank() ? HostAddress.of(tree.textValue()) : null;
            } else {
                var value = tree.get("value");
                var available = tree.get("available");
                if (value == null || !value.isTextual() || available == null || !available.isArray()) {
                    return null;
                }

                var l = new ArrayList<String>();
                for (JsonNode jsonNode : available) {
                    l.add(jsonNode.textValue());
                }
                return HostAddress.of(value.textValue(), l);
            }
        }
    }
}
