package org.ghidra.bridge.serialization;

import java.util.ArrayList;
import java.util.Comparator;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;

import ghidra.program.model.data.Array;
import ghidra.program.model.data.BitFieldDataType;
import ghidra.program.model.data.BuiltInDataType;
import ghidra.program.model.data.Composite;
import ghidra.program.model.data.DataType;
import ghidra.program.model.data.DataTypeComponent;
import ghidra.program.model.data.Enum;
import ghidra.program.model.data.FunctionDefinition;
import ghidra.program.model.data.ParameterDefinition;
import ghidra.program.model.data.Pointer;
import ghidra.program.model.data.Structure;
import ghidra.program.model.data.TypeDef;
import ghidra.program.model.data.Union;

public final class ProgramTypes {
    private final Map<String, String> ids = new LinkedHashMap<>();
    private final Map<String, DataType> values = new LinkedHashMap<>();

    public String id(DataType dataType) {
        if (dataType == null) {
            return null;
        }
        String key = key(dataType);
        String id = ids.get(key);
        if (id != null) {
            return id;
        }
        id = "type:" + idText(key);
        ids.put(key, id);
        values.put(key, dataType);
        collectNested(dataType);
        return id;
    }

    public List<Map<String, Object>> records() {
        List<String> keys = new ArrayList<>(values.keySet());
        keys.sort(String::compareTo);
        List<Map<String, Object>> records = new ArrayList<>();
        for (String key : keys) {
            records.add(record(ids.get(key), values.get(key)));
        }
        return records;
    }

    private void collectNested(DataType dataType) {
        if (dataType instanceof Pointer pointer) {
            id(pointer.getDataType());
        } else if (dataType instanceof Array array) {
            id(array.getDataType());
        } else if (dataType instanceof TypeDef typeDef) {
            id(typeDef.getDataType());
        } else if (dataType instanceof BitFieldDataType bitField) {
            id(bitField.getBaseDataType());
        } else if (dataType instanceof Composite composite) {
            for (DataTypeComponent component : composite.getDefinedComponents()) {
                id(component.getDataType());
            }
        } else if (dataType instanceof FunctionDefinition definition) {
            id(definition.getReturnType());
            for (ParameterDefinition parameter : definition.getArguments()) {
                id(parameter.getDataType());
            }
        }
    }

    private Map<String, Object> record(String id, DataType dataType) {
        Map<String, Object> value = new LinkedHashMap<>();
        String kind = kind(dataType);
        value.put("id", id);
        value.put("name", emptyToFallback(dataType.getName(), "unknown"));
        value.put("display_name", emptyToFallback(dataType.getDisplayName(), "unknown"));
        value.put("kind", kind);
        value.put("size", dataType.getLength());
        value.put("alignment", dataType.getAlignment());
        value.put(
            "category_path",
            dataType.getCategoryPath() == null ? null : emptyToNull(dataType.getCategoryPath().getPath())
        );

        switch (kind) {
            case "pointer" -> value.put("pointee_type_id", id(((Pointer) dataType).getDataType()));
            case "array" -> {
                Array array = (Array) dataType;
                value.put("element_type_id", id(array.getDataType()));
                value.put("element_count", array.getNumElements());
                value.put("element_size", array.getElementLength());
            }
            case "structure", "union" -> value.put("components", components((Composite) dataType));
            case "enum" -> {
                Enum enumType = (Enum) dataType;
                value.put("signed", enumType.isSigned());
                value.put("values", enumValues(enumType));
            }
            case "typedef" -> value.put("base_type_id", id(((TypeDef) dataType).getDataType()));
            case "function_definition" -> value.put(
                "signature",
                functionDefinitionSignature((FunctionDefinition) dataType)
            );
            case "bitfield" -> {
                BitFieldDataType bitField = (BitFieldDataType) dataType;
                value.put("base_type_id", id(bitField.getBaseDataType()));
                value.put("bit_size", bitField.getBitSize());
                value.put("bit_offset", bitField.getBitOffset());
                value.put("storage_size", bitField.getStorageSize());
            }
            default -> {
            }
        }
        return value;
    }

    private List<Map<String, Object>> components(Composite composite) {
        List<Map<String, Object>> records = new ArrayList<>();
        for (DataTypeComponent component : composite.getDefinedComponents()) {
            DataType dataType = component.getDataType();
            Map<String, Object> value = new LinkedHashMap<>();
            value.put("ordinal", component.getOrdinal());
            value.put("name", componentName(component));
            value.put("offset", component.getOffset());
            value.put("length", component.getLength());
            value.put("type_id", id(dataType));
            value.put("comment", emptyToNull(component.getComment()));
            if (dataType instanceof BitFieldDataType bitField) {
                value.put("bit_size", bitField.getBitSize());
                value.put("bit_offset", bitField.getBitOffset());
            } else {
                value.put("bit_size", null);
                value.put("bit_offset", null);
            }
            records.add(value);
        }
        records.sort(Comparator.comparing(record -> (Integer) record.get("ordinal")));
        return records;
    }

    private List<Map<String, Object>> enumValues(Enum enumType) {
        List<Map<String, Object>> records = new ArrayList<>();
        for (long value : enumType.getValues()) {
            for (String name : enumType.getNames(value)) {
                Map<String, Object> record = new LinkedHashMap<>();
                record.put("name", name);
                record.put("value", Long.toString(value));
                record.put("comment", emptyToNull(enumType.getComment(name)));
                records.add(record);
            }
        }
        records.sort((left, right) -> {
            int byValue = Long.compare(
                Long.parseLong((String) left.get("value")),
                Long.parseLong((String) right.get("value"))
            );
            if (byValue != 0) {
                return byValue;
            }
            return ((String) left.get("name")).compareTo((String) right.get("name"));
        });
        return records;
    }

    private Map<String, Object> functionDefinitionSignature(FunctionDefinition definition) {
        Map<String, Object> value = new LinkedHashMap<>();
        value.put("display", definition.getPrototypeString());
        value.put("calling_convention", emptyToFallback(definition.getCallingConventionName(), "unknown"));
        value.put("return_type_id", id(definition.getReturnType()));
        List<Map<String, Object>> parameters = new ArrayList<>();
        for (ParameterDefinition parameter : definition.getArguments()) {
            Map<String, Object> record = new LinkedHashMap<>();
            record.put("ordinal", parameter.getOrdinal());
            record.put("name", emptyToFallback(parameter.getName(), "param_" + parameter.getOrdinal()));
            record.put("type_id", id(parameter.getDataType()));
            record.put("storage", "unknown");
            parameters.add(record);
        }
        value.put("parameters", parameters);
        value.put("varargs", definition.hasVarArgs());
        value.put("no_return", definition.hasNoReturn());
        return value;
    }

    private static String key(DataType dataType) {
        return String.join(
            "\u001f",
            emptyToFallback(dataType.getName(), "unknown"),
            emptyToFallback(dataType.getDisplayName(), "unknown"),
            Integer.toString(dataType.getLength()),
            Integer.toString(dataType.getAlignment()),
            dataType.getCategoryPath() == null ? "" : dataType.getCategoryPath().getPath(),
            kind(dataType)
        );
    }

    private static String kind(DataType dataType) {
        if (dataType instanceof BitFieldDataType) {
            return "bitfield";
        }
        if (dataType instanceof Pointer) {
            return "pointer";
        }
        if (dataType instanceof Array) {
            return "array";
        }
        if (dataType instanceof Structure) {
            return "structure";
        }
        if (dataType instanceof Union) {
            return "union";
        }
        if (dataType instanceof Enum) {
            return "enum";
        }
        if (dataType instanceof TypeDef) {
            return "typedef";
        }
        if (dataType instanceof FunctionDefinition) {
            return "function_definition";
        }
        if (dataType instanceof BuiltInDataType) {
            return "builtin";
        }
        return "unknown";
    }

    private static String componentName(DataTypeComponent component) {
        String name = emptyToNull(component.getFieldName());
        if (name != null) {
            return name;
        }
        name = emptyToNull(component.getDefaultFieldName());
        return name == null ? "field_" + component.getOrdinal() : name;
    }

    public static String idText(String value) {
        return value.replace("%", "%25").replace(":", "%3a").replace("\u001f", "%1f");
    }

    private static String emptyToFallback(String value, String fallback) {
        return value == null || value.isBlank() ? fallback : value;
    }

    private static String emptyToNull(String value) {
        return value == null || value.isBlank() ? null : value;
    }
}
