# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.

import json
import time
from dataclasses import dataclass
from typing import Dict, List, Optional

from pypaimon.common.options.core_options import CoreOptions
from pypaimon.common.file_io import FileIO
from pypaimon.common.json_util import json_field
from pypaimon.schema.data_types import DataField
from pypaimon.schema.schema import Schema


@dataclass
class TableSchema:
    PAIMON_07_VERSION = 1
    PAIMON_08_VERSION = 2
    CURRENT_VERSION = 3

    FIELD_VERSION = "version"
    FIELD_ID = "id"
    FIELD_FIELDS = "fields"
    FIELD_HIGHEST_FIELD_ID = "highestFieldId"
    FIELD_PARTITION_KEYS = "partitionKeys"
    FIELD_PRIMARY_KEYS = "primaryKeys"
    FIELD_OPTIONS = "options"
    FIELD_COMMENT = "comment"
    FIELD_TIME_MILLIS = "timeMillis"

    version: int = json_field(FIELD_VERSION, default=CURRENT_VERSION)
    id: int = json_field(FIELD_ID, default=0)
    fields: List[DataField] = json_field(FIELD_FIELDS, default_factory=list)
    highest_field_id: int = json_field("highestFieldId", default=0)
    partition_keys: List[str] = json_field(
        FIELD_PARTITION_KEYS, default_factory=list)
    primary_keys: List[str] = json_field(
        FIELD_PRIMARY_KEYS, default_factory=list)
    options: Dict[str, str] = json_field(FIELD_OPTIONS, default_factory=dict)
    comment: Optional[str] = json_field(FIELD_COMMENT, default=None)
    time_millis: int = json_field("timeMillis", default_factory=lambda: int(time.time() * 1000))

    def cross_partition_update(self) -> bool:
        if not self.primary_keys or not self.partition_keys:
            return False

        # Check if primary keys contain all partition keys
        # Return True if they don't contain all (cross-partition update)
        return not all(pk in self.primary_keys for pk in self.partition_keys)

    @property
    def bucket_keys(self) -> List[str]:
        """Resolve the effective bucket-key column names.

        Resolution rule matches Java ``TableSchema.bucketKeys()``: prefer
        the explicit ``bucket-key`` option; otherwise fall back to primary
        keys with partition keys stripped (the same convention writers
        use).

        Validation is intentionally narrower than Java's
        ``originalBucketKeys()``: only ``unknown column name`` is checked
        here. Java additionally enforces ``bucket-key`` ⊄ partition keys,
        and (when primary keys are non-empty) ``bucket-key`` ⊆ primary
        keys, but it does so once at schema construction. Doing the same
        in a property would add per-read overhead and could surface
        errors on tables already in the catalog. The narrow check here
        is just enough to fail fast on the typo case.
        """
        configured = self.options.get(CoreOptions.BUCKET_KEY.key())
        if configured and configured.strip():
            keys = [k.strip() for k in configured.split(',') if k.strip()]
            field_names = {f.name for f in self.fields}
            missing = [k for k in keys if k not in field_names]
            if missing:
                raise ValueError(
                    "bucket-key references unknown columns: {}".format(missing))
            return keys
        return [pk for pk in self.primary_keys if pk not in self.partition_keys]

    @property
    def logical_bucket_key_fields(self) -> List[DataField]:
        """The ``DataField``s for ``bucket_keys``, in the order they were
        declared. Mirrors Java ``TableSchema.logicalBucketKeyType()``.
        """
        field_map = {f.name: f for f in self.fields}
        return [field_map[name] for name in self.bucket_keys]

    def to_schema(self) -> Schema:
        return Schema(
            fields=self.fields,
            partition_keys=self.partition_keys,
            primary_keys=self.primary_keys,
            options=self.options,
            comment=self.comment
        )

    @staticmethod
    def from_schema(schema_id: int, schema: Schema) -> "TableSchema":
        fields: List[DataField] = schema.fields
        partition_keys: List[str] = schema.partition_keys
        primary_keys: List[str] = schema.primary_keys
        options: Dict[str, str] = schema.options
        highest_field_id: int = max((field.id for field in fields), default=0)

        return TableSchema(
            TableSchema.CURRENT_VERSION,
            schema_id,
            fields,
            highest_field_id,
            partition_keys,
            primary_keys,
            options,
            schema.comment
        )

    @staticmethod
    def from_path(file_io: FileIO, schema_path: str):
        try:
            json_str = file_io.read_file_utf8(schema_path)
            return TableSchema.from_json(json_str)
        except FileNotFoundError as e:
            raise RuntimeError(f"Schema file not found: {schema_path}") from e
        except Exception as e:
            raise RuntimeError(f"Failed to read schema from {schema_path}") from e

    @staticmethod
    def from_json(json_str: str):
        try:
            data = json.loads(json_str)
            version = data.get(TableSchema.FIELD_VERSION, TableSchema.PAIMON_07_VERSION)
            fields = [DataField.from_dict(field) for field in data[TableSchema.FIELD_FIELDS]]
            options = data[TableSchema.FIELD_OPTIONS]
            if version <= TableSchema.PAIMON_07_VERSION and CoreOptions.BUCKET.key() not in options:
                options[CoreOptions.BUCKET.key()] = "1"
            if version <= TableSchema.PAIMON_08_VERSION and CoreOptions.FILE_FORMAT.key() not in options:
                options[CoreOptions.FILE_FORMAT.key()] = "orc"

            return TableSchema(
                version=version,
                id=data[TableSchema.FIELD_ID],
                fields=fields,
                highest_field_id=data[TableSchema.FIELD_HIGHEST_FIELD_ID],
                partition_keys=data[TableSchema.FIELD_PARTITION_KEYS],
                primary_keys=data[TableSchema.FIELD_PRIMARY_KEYS],
                options=options or {},
                comment=data.get(TableSchema.FIELD_COMMENT),
                time_millis=data.get(TableSchema.FIELD_TIME_MILLIS)
            )
        except json.JSONDecodeError as e:
            raise RuntimeError(f"Invalid JSON format: {json_str}") from e
        except KeyError as e:
            raise RuntimeError(f"Missing required field in schema JSON: {e}") from e
        except Exception as e:
            raise RuntimeError(f"Failed to parse schema from JSON: {e}") from e

    def copy(self, new_options: Optional[Dict[str, str]] = None) -> "TableSchema":
        return TableSchema(
            version=self.version,
            id=self.id,
            fields=self.fields,
            highest_field_id=self.highest_field_id,
            partition_keys=self.partition_keys,
            primary_keys=self.primary_keys,
            options=new_options,
            comment=self.comment,
            time_millis=self.time_millis
        )
