import pyarrow as pa
from pyarrow import Schema as PaSchema
from pyiceberg.schema import Schema as IcebergSchema
from pyiceberg.types import (
    NestedField,
    BooleanType,
    DateType,
    DoubleType,
    FloatType,
    IntegerType,
    LongType,
    StringType,
    TimestampType,
)
from typing import List

from pyiceberg.types import MapType

import pyarrow as pa
import json
from typing import List

def add_identifier_metadata_to_pyarrow_schema(
    schema: pa.Schema, identifier_field_names: List[str]
) -> pa.Schema:
    """
    Returns a new pyarrow.Schema with 'identifier_field_ids' metadata embedded.

    Parameters:
        schema (pa.Schema): Original schema
        identifier_field_names (List[str]): List of field names to mark as identifiers

    Returns:
        pa.Schema: New schema with metadata including 'identifier_field_ids'
    """
    field_ids = []
    for name in identifier_field_names:
        index = schema.get_field_index(name)
        if index == -1:
            raise ValueError(f"Field '{name}' not found in schema")
        field_ids.append(index)

    # Convert to JSON and encode as bytes
    identifier_meta = json.dumps(field_ids).encode("utf-8")

    # Copy or create metadata dictionary
    metadata = schema.metadata or {}
    metadata = dict(metadata)
    metadata[b"identifier_field_ids"] = identifier_meta

    return schema.with_metadata(metadata)




def pyarrow_to_iceberg_schema(pa_schema: PaSchema, identifier_field_names: List[str]) -> IcebergSchema:
    """
    Convert a PyArrow schema to a PyIceberg schema with unique field IDs and identifier fields.

    Parameters:
        pa_schema (pyarrow.Schema): The input PyArrow schema.
        identifier_field_names (List[str]): Field names to be used as identifier fields.

    Returns:
        pyiceberg.schema.Schema: A valid Iceberg schema with field IDs and identifier_field_ids.
    """

    def arrow_to_iceberg_type(field: pa.Field):
        pa_type = field.type

        if pa.types.is_int32(pa_type):
            return IntegerType()
        elif pa.types.is_int64(pa_type):
            return LongType()
        elif pa.types.is_float32(pa_type):
            return FloatType()
        elif pa.types.is_float64(pa_type):
            return DoubleType()
        elif pa.types.is_string(pa_type):
            return StringType()
        elif pa.types.is_boolean(pa_type):
            return BooleanType()
        elif pa.types.is_date(pa_type):
            return DateType()
        elif pa.types.is_timestamp(pa_type):
            return TimestampType()
        elif pa.types.is_map(pa_type):
            key_type = arrow_to_iceberg_type(pa.field("key", pa_type.key_type))
            item_type = arrow_to_iceberg_type(pa.field("value", pa_type.item_type))
            value_required = not pa_type.value_field.nullable  # NOW this is valid
            return MapType(
                key_type=key_type,
                value_type=item_type,
                value_required=value_required,
            )
        else:
            raise ValueError(f"Unsupported PyArrow type: {pa_type}")

    nested_fields = []
    identifier_field_ids = []
    field_id = 1  # Iceberg field IDs start at 1

    for field in pa_schema:
        iceberg_type = arrow_to_iceberg_type(field)
        nested_field = NestedField(
            field_id=field_id,
            name=field.name,
            field_type=iceberg_type,
            required=not field.nullable,
        )
        nested_fields.append(nested_field)
        if field.name in identifier_field_names:
            identifier_field_ids.append(field_id)
        field_id += 1

    return IcebergSchema(*nested_fields, identifier_field_ids=identifier_field_ids)