import copy
import json
from typing import Dict, Any, Union, List
from avrotize.common import build_tree_hash_list, group_by_hash, is_generic_json_type, NodeHashReference
from functools import reduce
import jsonpath_ng 

class AvroToJsonSchemaConverter:
    
    def __init__(self, naming_mode: str = 'snake') -> None:
        self.naming_mode = naming_mode
        self.defined_types: Dict[str, Any] = {}
        self.common_namespace = ''

    def find_common_namespace(self, namespaces: List[str]) -> str:
        """
        Find the common namespace prefix from a list of namespaces.
        """
        if not namespaces:
            return ''

        def common_prefix(a, b):
            prefix = ''
            for a_char, b_char in zip(a.split('.'), b.split('.')):
                if a_char == b_char:
                    prefix += a_char + '.'
                else:
                    break
            return prefix.rstrip('.')

        return reduce(common_prefix, namespaces)

    def update_common_namespace(self, namespace: str) -> None:
        """
        Update the common namespace based on the provided namespace.
        """
        if not self.common_namespace:
            self.common_namespace = namespace
        else:
            self.common_namespace = self.find_common_namespace([self.common_namespace, namespace])

    def get_definition_ref(self, name: str) -> str:
        """
        Construct the reference string based on the namespace and name.
        """
        
        if '.' in name:
            namespace, name = name.rsplit('.', 1)
        else:
            namespace = self.common_namespace
            
        if not self.common_namespace:
            return f"#/definitions/{name}"

        # Remove the common namespace and replace '.' with '/'
        namespace_suffix = namespace[len(self.common_namespace):].lstrip('.')
        path = namespace_suffix.replace('.', '/') if namespace_suffix else ''
        ref = f"#/definitions/{path}/{name}" if path else f"#/definitions/{name}"
        return ref
    
    def get_qualified_name(self, avro_type: Dict[str, Any]) -> str:
        """
        Construct the qualified name based on the namespace and name.
        """
        return avro_type['name'] if 'namespace' not in avro_type else f"{avro_type['namespace']}.{avro_type['name']}"

    def avro_primitive_to_json_type(self, avro_type: Union[str, Dict[str, Any]]) -> Dict[str, Any]:
        """
        Map Avro primitive types to JSON types with appropriate format annotations.
        Handles both standard Avro logical types and Avrotize schema extensions.
        """
        json_type = {}
        if isinstance(avro_type, dict):
            # Check for logical type before unwrapping the base type
            logical_type = avro_type.get('logicalType')
            base_type = avro_type.get('type', avro_type)
            
            # Handle logical types based on their base type
            if logical_type and isinstance(base_type, str):
                # Standard Avro logical types on int/long
                if base_type == 'int' and logical_type == 'date':
                    # Standard Avro: int with date logicalType represents days since epoch
                    json_type['type'] = 'integer'
                    json_type['format'] = 'int32'
                    return json_type
                elif base_type == 'int' and logical_type in ['time-millis']:
                    # Standard Avro: int with time-millis represents milliseconds since midnight
                    json_type['type'] = 'integer'
                    json_type['format'] = 'int32'
                    return json_type
                elif base_type == 'long' and logical_type in ['time-micros']:
                    # Standard Avro: long with time-micros represents microseconds since midnight
                    json_type['type'] = 'integer'
                    json_type['format'] = 'int64'
                    return json_type
                elif base_type == 'long' and logical_type in ['timestamp-millis', 'timestamp-micros']:
                    # Standard Avro: long with timestamp represents milliseconds/microseconds since epoch
                    json_type['type'] = 'integer'
                    json_type['format'] = 'int64'
                    return json_type
                # Avrotize schema extensions: string-based logical types
                elif base_type == 'string' and logical_type == 'date':
                    # Avrotize extension: string with date logicalType
                    json_type['type'] = 'string'
                    json_type['format'] = 'date'
                    return json_type
                elif base_type == 'string' and logical_type in ['timestamp-millis', 'timestamp-micros', 'datetime']:
                    # Avrotize extension: string with datetime logicalType
                    json_type['type'] = 'string'
                    json_type['format'] = 'date-time'
                    return json_type
                elif base_type == 'string' and logical_type in ['time-millis', 'time-micros', 'time']:
                    # Avrotize extension: string with time logicalType
                    json_type['type'] = 'string'
                    json_type['format'] = 'time'
                    return json_type
                elif logical_type == 'decimal':
                    json_type['type'] = 'number'
                    return json_type
                elif logical_type == 'uuid':
                    json_type['type'] = 'string'
                    json_type['format'] = 'uuid'
                    return json_type
            
            # If base_type is still a dict, recurse
            if isinstance(base_type, dict):
                if 'logicalType' in base_type:
                    return self.avro_primitive_to_json_type(base_type)
                else:
                    raise ValueError(f"Avro schema contains unexpected construct {avro_type}")
            
            # No logical type or unhandled combination, process base type
            return self.avro_primitive_to_json_type(base_type)
            
        mapping = {
            'null': {'type': 'null'},
            'boolean': {'type': 'boolean'},
            'int': {'type': 'integer', 'format': 'int32'},
            'long': {'type': 'integer', 'format': 'int64'},
            'float': {'type': 'number', 'format': 'float'},
            'double': {'type': 'number', 'format': 'double'},
            'bytes': {'type': 'string', 'contentEncoding': 'base64'},
            'string': {'type': 'string'},
            'fixed': {'type': 'string'}  # Could specify length in a format or a separate attribute
        }
        type_ref = mapping.get(avro_type, '')  # Defaulting to string type for any unknown types
        if not type_ref:
            raise ValueError(f"Avro schema contains unexpected type {avro_type}")
        return type_ref


    def convert_name(self, name: str) -> str:
        """
        Convert names according to the specified naming mode.
        """
        if self.naming_mode == 'snake':
            return self.to_snake_case(name)
        elif self.naming_mode == 'camel':
            return self.to_camel_case(name)
        elif self.naming_mode == 'pascal':
            return self.to_pascal_case(name)
        return name

    @staticmethod
    def to_snake_case(name: str) -> str:
        return ''.join(['_'+c.lower() if c.isupper() else c for c in name]).lstrip('_')

    @staticmethod
    def to_camel_case(name: str) -> str:
        return ''.join(word.capitalize() if i else word for i, word in enumerate(name.split('_')))

    @staticmethod
    def to_pascal_case(name: str) -> str:
        return ''.join(word.capitalize() for word in name.split('_'))

    def is_nullable(self, avro_type: Union[str, Dict[str, Any]]) -> bool:
        """
        Check if a given Avro type is nullable.
        """
        if isinstance(avro_type, list):
            return 'null' in avro_type
        return avro_type == 'null'
    
    def handle_type_union(self, types: List[Union[str, Dict[str, Any]]]) -> Dict[str, Any] | List[Dict[str, Any]| str] | str:
        """
        Handle Avro type unions, returning a JSON schema that validates against any of the types.
        """
        non_null_types = [t for t in types if t != 'null']
        if len(non_null_types) == 1:
            # Single non-null type
            return self.parse_avro_schema(non_null_types[0])
        else:
            # Multiple non-null types
            union_types = [self.convert_reference(t) if isinstance(t,str) and t in self.defined_types else self.avro_primitive_to_json_type(t)
                            if isinstance(t, str) else self.parse_avro_schema(t)
                                for t in non_null_types]
            return {
                'oneOf': union_types
            }

    def parse_avro_schema(self, avro_schema: Dict[str, Any] | List[Dict[str, Any]| str] | str, is_root = False) -> Dict[str, Any] | List[Dict[str, Any]| str] | str:
        """
        Parse an Avro schema structure and return the corresponding JSON schema.
        """
        if isinstance(avro_schema, list):
            # Type union
            union = self.handle_type_union(avro_schema)
            if is_root:
                # all the definitions go into 'definitions'
                return {
                    "$schema": "http://json-schema.org/draft-07/schema#"
                }
            if is_generic_json_type(union):
                return { "type": "object" }
            else:
                return union
        elif isinstance(avro_schema, dict):
            if 'namespace' in avro_schema:
                namespace = avro_schema['namespace']
                self.update_common_namespace(namespace)
            if avro_schema['type'] == 'record':
                return self.convert_record(avro_schema, is_root)
            elif avro_schema['type'] == 'enum':
                return self.convert_enum(avro_schema, is_root)
            elif avro_schema['type'] == 'fixed':
                return self.convert_fixed(avro_schema, is_root)
            elif avro_schema['type'] == 'array':
                return self.convert_array(avro_schema)
            elif avro_schema['type'] == 'map':
                return self.convert_map(avro_schema)
            elif avro_schema['type'] in self.defined_types:
                # Type reference
                return self.convert_reference(avro_schema)
            else:
                # Nested type or a direct type definition
                return self.parse_avro_schema(avro_schema['type'])
        elif isinstance(avro_schema, str):
            # Primitive type or a reference to a defined type
            if avro_schema in self.defined_types:
                return self.convert_reference(avro_schema)
            elif '.' in avro_schema:
                raise ValueError(f"Unknown type reference {avro_schema}")
            else:
                return self.avro_primitive_to_json_type(avro_schema)

    def convert_reference(self, avro_schema: Dict[str, Any] | str) -> Dict[str, Any]:
        """
        Convert a reference to a defined type to a JSON schema object with a reference to the definition.
        """
        key = avro_schema['type'] if isinstance(avro_schema, dict) else avro_schema
        json_type = self.defined_types[key]
        if 'enum' in json_type:
            return copy.deepcopy(json_type)
        else:
            return {"$ref": self.get_definition_ref(key)}
    
    def convert_record(self, avro_schema: Dict[str, Any], is_root=False) -> Dict[str, Any]:
        """
        Convert an Avro record type to a JSON schema object, handling nested types and type definitions.
        """
        record_name = self.convert_name(avro_schema['name'])
        properties = {}
        required = []
        
        json_schema: Dict[str, Any] = {
            "type": "object",
            "title": record_name
        }        
        if not is_root:
            self.defined_types[self.get_qualified_name(avro_schema)] = json_schema
        
        for field in avro_schema['fields']:
            field_name = self.convert_name(field['name'])
            prop = self.parse_avro_schema(field['type'])
            if 'doc' in field:
                if isinstance(prop, dict):
                    prop['description'] = field['doc']
                elif isinstance(prop, list) or isinstance(prop, str):
                    prop = {
                        'allOf': [
                            prop,
                            {'description': field['doc']}
                        ]}
            properties[field_name] = prop
            if not self.is_nullable(field['type']):
                required.append(field_name)

        if 'doc' in avro_schema:
            json_schema['description'] = avro_schema['doc']
        if properties:
            json_schema['properties'] = properties
        
        if required:
            json_schema['required'] = required
            
        if not is_root:
            return {"$ref": self.get_definition_ref(self.get_qualified_name(avro_schema))}
        return json_schema

    def convert_enum(self, avro_schema: Dict[str, Any], is_root=False) -> Dict[str, Any]:
        """
        Convert an Avro enum type to a JSON schema enum, adding the definition to the schema.
        """
        enum_name = self.convert_name(avro_schema['name'])
        json_schema = {
            "type": "string",
            "enum": avro_schema['symbols'],
            "title": enum_name
        }

        if 'doc' in avro_schema:
            json_schema['description'] = avro_schema['doc']

        # Add to defined types
        if not is_root:
            self.defined_types[self.get_qualified_name(avro_schema)] = json_schema
        return json_schema

    def convert_fixed(self, avro_schema: Dict[str, Any], is_root=False) -> Dict[str, Any]:
        """
        Convert an Avro fixed type to a JSON schema string with length constraints.
        Fixed types are represented as strings with base16 (hex) encoding and 
        minLength and maxLength constraints based on the size.
        """
        fixed_name = self.convert_name(avro_schema['name'])
        size = avro_schema['size']
        # Fixed types in JSON are represented as hex strings, so length is 2 * size
        hex_length = size * 2
        
        json_schema = {
            "type": "string",
            "contentEncoding": "base16",
            "minLength": hex_length,
            "maxLength": hex_length,
            "title": fixed_name
        }

        if 'doc' in avro_schema:
            json_schema['description'] = avro_schema['doc']

        # Add to defined types
        if not is_root:
            self.defined_types[self.get_qualified_name(avro_schema)] = json_schema
            return {"$ref": self.get_definition_ref(self.get_qualified_name(avro_schema))}
        return json_schema

    def convert_array(self, avro_schema: Dict[str, Any]) -> Dict[str, Any]:
        """
        Convert an Avro array type to a JSON schema array.
        """
        return {
            "type": "array",
            "items": self.parse_avro_schema(avro_schema['items'])
        }

    def convert_map(self, avro_schema: Dict[str, Any]) -> Dict[str, Any]:
        """
        Convert an Avro map type to a JSON schema object with additionalProperties.
        """
        return {
            "type": "object",
            "additionalProperties": self.parse_avro_schema(avro_schema['values'])
        }

    def convert(self, avro_schema: Dict[str, Any] | List[Dict[str, Any]| str] | str) -> Dict[str, Any] | List[Dict[str, Any]| str] | str:
        """
        Convert the root Avro schema to a JSON schema.
        """
        json_schema: Dict[str, Any] | List[Dict[str, Any]| str] | str = self.parse_avro_schema(avro_schema, is_root = True)
        
        if self.defined_types and isinstance(json_schema, dict):
            for name, definition in self.defined_types.items():
                if isinstance(definition, dict) and 'enum' in definition:
                    # enums are inlined
                    continue
                current_level = json_schema.setdefault('definitions', {})
                if '.' in name:
                    definition_namespace, definition_name = name.rsplit('.',1)
                    if not self.common_namespace or (self.common_namespace and definition_namespace == self.common_namespace):
                        definition_namespace = ''
                    else:
                        definition_namespace = definition_namespace[len(self.common_namespace):].lstrip('.')
                    # Split the definition_namespace into path segments
                    path_segments = definition_namespace.split('.')
                    if definition_namespace and len(path_segments) > 0:
                        # Traverse through all but the last segment, creating nested dictionaries as needed
                        for segment in path_segments:
                            # If the segment does not exist, create a new dictionary at that level
                            if segment not in current_level:
                                current_level[segment] = {}
                            # Move deeper into the nested structure
                            current_level = current_level[segment]
                else:
                    definition_name = name
                current_level[definition_name] = copy.deepcopy(definition)
                
        return json_schema

def compact_tree(json_schema):
    shared_def_counter = 1
    ignored_hashes = []
    while True:
        thl = build_tree_hash_list(json_schema)
        ghl = group_by_hash(thl)
        if len(ghl) == 0:
            return
        # sort ghl by the count in of the first item in each group
        ghl = dict(sorted(ghl.items(), key=lambda item: -item[1][0].count))
        repeat = True
        while repeat:
            repeat = False
            first_group_key = next((key for key in ghl.keys() if key not in ignored_hashes), None)
            if first_group_key is None:
                return
            ghl_top_item_entries = ghl[first_group_key]
            # sort the items by the shortest .path value
            ghl_top_item_entries = sorted(ghl_top_item_entries, key=lambda item: len(item.path.split('.')))
            top_item_entry: NodeHashReference = ghl_top_item_entries[0]
            top_item_path_segments = top_item_entry.path.split('.')
            if top_item_path_segments[1] == 'definitions' and len(top_item_path_segments) == 3:
                # the top item sits right under definitions, we will merge into that one
                def_key = top_item_path_segments[2]
                ghl_top_item_entries.remove(top_item_entry)
            elif ((top_item_path_segments[-1] == 'options' and top_item_path_segments[-2] == 'properties' and len(top_item_path_segments) > 4) and 'oneOf' in top_item_entry.value):
                # the first case is likely a union we created in j2a that we had to create a top-level item for. We will undo that here.
                json_item = json_schema
                def_key = ''
                for seg in top_item_path_segments[1:-2]:
                    def_key += '/' + seg if def_key else seg
                    json_item = json_item[seg]
                json_item.clear()
                json_item.update(copy.deepcopy(top_item_entry.value))
                ghl_top_item_entries.remove(top_item_entry)    
            elif top_item_path_segments[-2] == 'properties' or top_item_path_segments[-1] == 'properties':
                # the top item is a property of an object, which means that we would create direct
                # links into that object and therefore we will drop that hash
                ignored_hashes.append(first_group_key)
                repeat = True
                continue
            else:
                # the second is indeed a proper type declaration, so we will use the first as the one all other occurrences refer to
                json_item = json_schema
                def_key = ''
                for seg in top_item_path_segments[1:]:
                    def_key += '/' + seg if def_key else seg
                ghl_top_item_entries.remove(top_item_entry)        
        
        
            for ghl_item in ghl_top_item_entries:
                node = ghl_item.value
                if isinstance(node,dict):
                    node.clear()
                    node.update({
                        '$ref': f"#/{def_key}"
                    })
            break



def convert_avro_to_json_schema(avro_schema_file: str, json_schema_file: str, naming_mode: str = 'default') -> None:
    """
    Convert an Avro schema file to a JSON schema file.

    :param avro_schema_file: The path to the input Avro schema file.
    :param json_schema_file: The path to the output JSON schema file.
    :param naming_mode: The naming mode for converting names ('snake', 'camel', 'pascal').
    """
    converter = AvroToJsonSchemaConverter(naming_mode)

    # Read the Avro schema file
    with open(avro_schema_file, 'r') as file:
        avro_schema = json.load(file)

    # Convert the Avro schema to JSON schema
    json_schema = converter.convert(avro_schema)
    
    compact_tree(json_schema)
    # Write the JSON schema to the output file
    with open(json_schema_file, 'w') as file:
        json.dump(json_schema, file, indent=4)

