"""Convert a JSON Structure schema to an Iceberg schema."""

import json
import sys
from typing import Dict, List, Any, Optional
import pyarrow as pa
from pyiceberg.schema import Schema, NestedField
from pyiceberg.types import (
    BooleanType,
    IntegerType,
    LongType,
    FloatType,
    DoubleType,
    StringType,
    BinaryType,
    DateType,
    TimestampType,
    DecimalType,
    FixedType,
    ListType,
    MapType,
    StructType,
    TimeType
)
from pyiceberg.io.pyarrow import PyArrowFileIO, schema_to_pyarrow

JsonNode = Dict[str, 'JsonNode'] | List['JsonNode'] | str | bool | int | None


class StructureToIcebergConverter:
    """Class to convert JSON Structure schema to Iceberg schema."""

    def __init__(self: 'StructureToIcebergConverter'):
        self.named_type_cache: Dict[str, JsonNode] = {}
        self.id_counter = 0
        self.definitions: Dict[str, Any] = {}
        self.schema_doc: Optional[Dict[str, Any]] = None

    def get_id(self) -> int:
        """Get a unique ID for a field."""
        self.id_counter += 1
        return self.id_counter
    
    def get_fullname(self, namespace: str, name: str) -> str:
        """Get the full name of a record type."""
        return f"{namespace}.{name}" if namespace else name

    def convert_structure_to_iceberg(self, structure_schema_path: str, structure_record_type: Optional[str], output_path: str, emit_cloudevents_columns: bool=False):
        """Convert a JSON Structure schema to an Iceberg schema."""
        schema_file = structure_schema_path
        if not schema_file:
            print("Please specify the JSON Structure schema file")
            sys.exit(1)
        with open(schema_file, "r", encoding="utf-8") as f:
            schema_json = f.read()

        # Parse the schema as a JSON object
        schema = json.loads(schema_json)
        self.schema_doc = schema
        
        # Handle definitions if present
        if "definitions" in schema:
            self.definitions = schema["definitions"]
        
        # For JSON Structure, we expect an object type at the top level
        if schema.get("type") != "object":
            # Check if we have a $ref at the top level
            if "$ref" in schema:
                ref = schema["$ref"]
                schema = self.resolve_ref(ref)
            elif structure_record_type and "definitions" in schema:
                # Look for the type in definitions
                if structure_record_type in schema["definitions"]:
                    schema = schema["definitions"][structure_record_type]
                else:
                    print(f"No record type {structure_record_type} found in the JSON Structure schema definitions")
                    sys.exit(1)
            else:
                print("Expected a JSON Structure schema with type 'object' at the top level")
                sys.exit(1)

        # Get the name and properties of the top-level object
        table_name = schema.get("name", "Table")
        properties = schema.get("properties", {})
        required = schema.get("required", [])

        # Create a list to store the iceberg schema
        iceberg_fields: List[NestedField] = []

        # Append the iceberg schema with the column names and types
        for prop_name, prop_schema in properties.items():
            is_required = prop_name in required
            column_type = self.convert_structure_type_to_iceberg_type(prop_schema)
            iceberg_fields.append(
                NestedField(
                    field_id=self.get_id(),
                    name=prop_name,
                    field_type=column_type,
                    required=is_required
                ))

        if emit_cloudevents_columns:
            iceberg_fields.extend([
                NestedField(field_id=self.get_id(),
                      name="___type", field_type=StringType(), required=False),
                NestedField(field_id=self.get_id(),
                      name="___source", field_type=StringType(), required=False),
                NestedField(field_id=self.get_id(),
                      name="___id", field_type=StringType(), required=False),
                NestedField(field_id=self.get_id(),
                      name="___time", field_type=TimestampType(), required=False),
                NestedField(field_id=self.get_id(),
                      name="___subject", field_type=StringType(), required=False)
            ])

        iceberg_schema = Schema(*iceberg_fields)
        arrow_schema = schema_to_pyarrow(iceberg_schema)
        print(f"Iceberg schema created: {arrow_schema}")

        # Write to Iceberg table (for demonstration, using local file system)
        file_io = PyArrowFileIO()
        output_file = file_io.new_output("file://"+output_path)
        with output_file.create(overwrite=True) as f:
            pa.output_stream(f).write(arrow_schema.serialize().to_pybytes())

    def resolve_ref(self, ref: str) -> Dict[str, Any]:
        """Resolve a $ref reference."""
        if not ref.startswith("#/"):
            raise ValueError(f"Only local references are supported, got: {ref}")
        
        parts = ref[2:].split("/")
        current = self.schema_doc
        
        for part in parts:
            if isinstance(current, dict) and part in current:
                current = current[part]
            else:
                raise ValueError(f"Could not resolve reference: {ref}")
        
        return current

    def convert_structure_type_to_iceberg_type(self, structure_type):
        """Convert a JSON Structure type to an Iceberg type."""
        # Handle $ref
        if isinstance(structure_type, dict) and "$ref" in structure_type:
            ref = structure_type["$ref"]
            resolved = self.resolve_ref(ref)
            return self.convert_structure_type_to_iceberg_type(resolved)
        
        # Handle array of types (e.g., ["string", "null"] for nullable types)
        if isinstance(structure_type, list):
            # Filter out null from the list
            non_null_types = [t for t in structure_type if t != "null"]
            if len(non_null_types) == 1:
                # Nullable type - just use the non-null type (Iceberg handles optionality with required flag)
                return self.convert_structure_type_to_iceberg_type(non_null_types[0])
            elif len(non_null_types) > 1:
                # Union of multiple non-null types - create a struct with alternatives
                fields = []
                for i, choice in enumerate(non_null_types):
                    choice_type = self.convert_structure_type_to_iceberg_type(choice)
                    fields.append(NestedField(
                        field_id=self.get_id(),
                        name=f"option_{i}",
                        field_type=choice_type,
                        required=False
                    ))
                return StructType(*fields)
            else:
                # Only null - return string as fallback
                return StringType()
        
        # Handle dictionary with type field
        if isinstance(structure_type, dict):
            type_name = structure_type.get("type")
            
            # Handle type being an array
            if isinstance(type_name, list):
                # This is like {"type": ["string", "null"]}
                return self.convert_structure_type_to_iceberg_type(type_name)
            
            # Handle array type
            if type_name == "array":
                items = structure_type.get("items", {"type": "string"})
                return ListType(
                    element_id=self.get_id(),
                    element_type=self.convert_structure_type_to_iceberg_type(items),
                    element_required=True
                )
            
            # Handle set type (treated as array in Iceberg)
            elif type_name == "set":
                items = structure_type.get("items", {"type": "string"})
                return ListType(
                    element_id=self.get_id(),
                    element_type=self.convert_structure_type_to_iceberg_type(items),
                    element_required=True
                )
            
            # Handle map type
            elif type_name == "map":
                values = structure_type.get("values", {"type": "string"})
                return MapType(
                    key_id=self.get_id(),
                    key_type=StringType(),
                    value_id=self.get_id(),
                    value_type=self.convert_structure_type_to_iceberg_type(values),
                    value_required=True
                )
            
            # Handle tuple type (treated as struct with indexed fields)
            elif type_name == "tuple":
                items = structure_type.get("items", [])
                fields = []
                for i, item in enumerate(items):
                    fields.append(NestedField(
                        field_id=self.get_id(),
                        name=f"field_{i}",
                        field_type=self.convert_structure_type_to_iceberg_type(item),
                        required=True
                    ))
                return StructType(*fields)
            
            # Handle object type
            elif type_name == "object":
                properties = structure_type.get("properties", {})
                required = structure_type.get("required", [])
                fields = []
                
                # Handle $extends if present
                if "$extends" in structure_type:
                    extends_ref = structure_type["$extends"]
                    base_schema = self.resolve_ref(extends_ref)
                    base_properties = base_schema.get("properties", {})
                    base_required = base_schema.get("required", [])
                    
                    # Add base properties first
                    for prop_name, prop_schema in base_properties.items():
                        is_required = prop_name in base_required
                        fields.append(NestedField(
                            field_id=self.get_id(),
                            name=prop_name,
                            field_type=self.convert_structure_type_to_iceberg_type(prop_schema),
                            required=is_required
                        ))
                
                # Add own properties
                for prop_name, prop_schema in properties.items():
                    is_required = prop_name in required
                    fields.append(NestedField(
                        field_id=self.get_id(),
                        name=prop_name,
                        field_type=self.convert_structure_type_to_iceberg_type(prop_schema),
                        required=is_required
                    ))
                
                return StructType(*fields)
            
            # Handle choice type (union)
            elif type_name == "choice":
                choices = structure_type.get("choices", [])
                if isinstance(choices, list):
                    # For inline choices, create a struct with alternatives
                    fields = []
                    for i, choice in enumerate(choices):
                        choice_type = self.convert_structure_type_to_iceberg_type(choice)
                        fields.append(NestedField(
                            field_id=self.get_id(),
                            name=f"option_{i}",
                            field_type=choice_type,
                            required=False
                        ))
                    return StructType(*fields)
                elif isinstance(choices, dict):
                    # For tagged choices, create a struct with named alternatives
                    fields = []
                    for choice_name, choice_schema in choices.items():
                        choice_type = self.convert_structure_type_to_iceberg_type(choice_schema)
                        fields.append(NestedField(
                            field_id=self.get_id(),
                            name=choice_name,
                            field_type=choice_type,
                            required=False
                        ))
                    return StructType(*fields)
                else:
                    return StringType()
            
            # Handle any type
            elif type_name == "any":
                return StringType()
            
            # Handle primitive types with annotations
            elif type_name:
                return self.map_iceberg_scalar_type(type_name, structure_type)
        
        # Handle string type name directly
        elif isinstance(structure_type, str):
            return self.map_iceberg_scalar_type(structure_type, {})

        return StringType()

    def map_iceberg_scalar_type(self, type_name: str, type_schema: Dict[str, Any]):
        """Map a JSON Structure scalar type to an Iceberg scalar type."""
        # Check for decimal with precision and scale
        if type_name == "decimal":
            precision = type_schema.get("precision", 38)
            scale = type_schema.get("scale", 18)
            return DecimalType(precision, scale)
        
        # Map other primitive types
        type_mapping = {
            'null': StringType(),  # Iceberg doesn't have a null type
            'boolean': BooleanType(),
            'string': StringType(),
            'int8': IntegerType(),  # Iceberg doesn't have byte type
            'uint8': IntegerType(),
            'int16': IntegerType(),  # Iceberg doesn't have short type
            'uint16': IntegerType(),
            'int32': IntegerType(),
            'uint32': LongType(),  # Use long for unsigned int32
            'int64': LongType(),
            'uint64': LongType(),  # Iceberg doesn't distinguish signed/unsigned
            'int128': StringType(),  # No native 128-bit support
            'uint128': StringType(),
            'integer': IntegerType(),  # Generic integer
            'number': DoubleType(),  # Generic number
            'float8': FloatType(),
            'float': FloatType(),
            'float32': FloatType(),
            'binary32': FloatType(),
            'double': DoubleType(),
            'float64': DoubleType(),
            'binary64': DoubleType(),
            'decimal': DecimalType(38, 18),
            'binary': BinaryType(),
            'bytes': BinaryType(),  # Binary data
            'date': DateType(),
            'time': TimeType(),
            'datetime': TimestampType(),
            'timestamp': TimestampType(),
            'duration': LongType(),  # Store as microseconds
            'uuid': StringType(),  # Store UUID as string
            'uri': StringType(),
            'jsonpointer': StringType(),
        }
        
        return type_mapping.get(type_name, StringType())


def convert_structure_to_iceberg(structure_schema_path, structure_record_type, output_path, emit_cloudevents_columns=False):
    """Convert a JSON Structure schema to an Iceberg schema."""
    converter = StructureToIcebergConverter()
    converter.convert_structure_to_iceberg(
        structure_schema_path, structure_record_type, output_path, emit_cloudevents_columns)
