# 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 logging
import threading
import time
from typing import Optional, Union

from cachetools import TTLCache

from pypaimon.api.rest_api import RESTApi
from pypaimon.api.rest_util import RESTUtil
from pypaimon.catalog.rest.rest_token import RESTToken
from pypaimon.common.file_io import FileIO
from pypaimon.filesystem.pyarrow_file_io import PyArrowFileIO
from pypaimon.api.auth.bearer import BearTokenAuthProvider
from pypaimon.api.auth.dlf_provider import DLFAuthProvider
from pypaimon.common.identifier import Identifier
from pypaimon.common.options import Options
from pypaimon.common.options.config import CatalogOptions, OssOptions
from pypaimon.common.uri_reader import UriReaderFactory


class RESTTokenFileIO(FileIO):
    """
    A FileIO to support getting token from REST Server.
    """

    _FILE_IO_CACHE_MAXSIZE = 1000
    _FILE_IO_CACHE_TTL = 36000  # 10 hours in seconds

    _FILE_IO_CACHE: TTLCache = None
    _FILE_IO_CACHE_LOCK = threading.Lock()

    _TOKEN_CACHE: dict = {}
    _TOKEN_LOCKS: dict = {}
    _TOKEN_LOCKS_LOCK = threading.Lock()

    @classmethod
    def _get_file_io_cache(cls) -> TTLCache:
        if cls._FILE_IO_CACHE is None:
            with cls._FILE_IO_CACHE_LOCK:
                if cls._FILE_IO_CACHE is None:
                    cls._FILE_IO_CACHE = TTLCache(
                        maxsize=cls._FILE_IO_CACHE_MAXSIZE,
                        ttl=cls._FILE_IO_CACHE_TTL
                    )
        return cls._FILE_IO_CACHE

    def __init__(self, identifier: Identifier, path: str,
                 catalog_options: Optional[Union[dict, Options]] = None):
        self.identifier = identifier
        self.path = path
        if catalog_options is None:
            self.catalog_options = None
        elif isinstance(catalog_options, dict):
            self.catalog_options = Options(catalog_options)
        else:
            # Assume it's already an Options object
            self.catalog_options = catalog_options
        self.properties = self.catalog_options or Options({})  # For compatibility with refresh_token()
        self.token: Optional[RESTToken] = None
        self.api_instance: Optional[RESTApi] = None
        self.log = logging.getLogger(__name__)
        self._uri_reader_factory_cache: Optional[UriReaderFactory] = None

    def __getstate__(self):
        state = self.__dict__.copy()
        # Remove non-serializable objects
        state.pop('api_instance', None)
        state.pop('_uri_reader_factory_cache', None)
        # token can be serialized, but we'll refresh it on deserialization
        return state

    def __setstate__(self, state):
        self.__dict__.update(state)
        self._uri_reader_factory_cache = None
        # api_instance will be recreated when needed
        self.api_instance = None

    def file_io(self) -> FileIO:
        self.try_to_refresh_token()

        if self.token is None:
            return FileIO.get(self.path, self.catalog_options or Options({}))

        cache_key = self.token
        cache = self._get_file_io_cache()

        file_io = cache.get(cache_key)
        if file_io is not None:
            return file_io

        with self._FILE_IO_CACHE_LOCK:
            self.try_to_refresh_token()

            if self.token is None:
                return FileIO.get(self.path, self.catalog_options or Options({}))

            cache_key = self.token
            cache = self._get_file_io_cache()
            file_io = cache.get(cache_key)
            if file_io is not None:
                return file_io

            merged_properties = RESTUtil.merge(
                self.catalog_options.to_map() if self.catalog_options else {},
                self.token.token
            )
            if self.catalog_options:
                dlf_oss_endpoint = self.catalog_options.get(CatalogOptions.DLF_OSS_ENDPOINT)
                if dlf_oss_endpoint and dlf_oss_endpoint.strip():
                    merged_properties[OssOptions.OSS_ENDPOINT.key()] = dlf_oss_endpoint
            merged_options = Options(merged_properties)

            file_io = PyArrowFileIO(self.path, merged_options)
            cache[cache_key] = file_io
            return file_io

    def _merge_token_with_catalog_options(self, token: dict) -> dict:
        """Merge token with catalog options, DLF OSS endpoint should override the standard OSS endpoint."""
        merged_token = dict(token)
        if self.catalog_options:
            dlf_oss_endpoint = self.catalog_options.get(CatalogOptions.DLF_OSS_ENDPOINT)
            if dlf_oss_endpoint and dlf_oss_endpoint.strip():
                merged_token[OssOptions.OSS_ENDPOINT.key()] = dlf_oss_endpoint
        return merged_token

    def new_input_stream(self, path: str):
        return self.file_io().new_input_stream(path)

    def new_output_stream(self, path: str):
        return self.file_io().new_output_stream(path)

    def get_file_status(self, path: str):
        return self.file_io().get_file_status(path)

    def list_status(self, path: str):
        return self.file_io().list_status(path)

    def exists(self, path: str) -> bool:
        return self.file_io().exists(path)

    def delete(self, path: str, recursive: bool = False) -> bool:
        return self.file_io().delete(path, recursive)

    def mkdirs(self, path: str) -> bool:
        return self.file_io().mkdirs(path)

    def rename(self, src: str, dst: str) -> bool:
        return self.file_io().rename(src, dst)

    def copy_file(self, source_path: str, target_path: str, overwrite: bool = False):
        return self.file_io().copy_file(source_path, target_path, overwrite)

    def to_filesystem_path(self, path: str) -> str:
        return self.file_io().to_filesystem_path(path)

    def try_to_write_atomic(self, path: str, content: str) -> bool:
        return self.file_io().try_to_write_atomic(path, content)

    def write_parquet(self, path: str, data, compression: str = 'zstd',
                      zstd_level: int = 1, **kwargs):
        return self.file_io().write_parquet(path, data, compression, zstd_level, **kwargs)

    def write_orc(self, path: str, data, compression: str = 'zstd',
                  zstd_level: int = 1, **kwargs):
        return self.file_io().write_orc(path, data, compression, zstd_level, **kwargs)

    def write_avro(self, path: str, data, avro_schema=None,
                   compression: str = 'zstd', zstd_level: int = 1, **kwargs):
        return self.file_io().write_avro(path, data, avro_schema, compression, zstd_level, **kwargs)

    def write_lance(self, path: str, data, **kwargs):
        return self.file_io().write_lance(path, data, **kwargs)

    def write_blob(self, path: str, data, **kwargs):
        return self.file_io().write_blob(path, data, **kwargs)

    def write_row(self, path: str, data, fields=None, zstd_level: int = 1, **kwargs):
        return self.file_io().write_row(path, data, fields, zstd_level, **kwargs)

    @property
    def uri_reader_factory(self):
        if self._uri_reader_factory_cache is None:
            catalog_options = self.catalog_options or Options({})
            self._uri_reader_factory_cache = UriReaderFactory(catalog_options)

        return self._uri_reader_factory_cache

    @property
    def filesystem(self):
        return self.file_io().filesystem

    def _build_cache_key(self) -> str:
        """Build cache key with user identity to isolate tokens across different credentials."""
        return f"{self._get_user_identity()}:{self.identifier}"

    def _get_user_identity(self) -> str:
        """Get user identity for token cache isolation.

        For DLF auth (AK/SK, STS, ECS Role): returns access_key_id via get_token().
        For Bear Token auth: returns the bearer token string.
        For unknown auth: returns 'anonymous'.
        """
        if self.api_instance is None:
            self.api_instance = RESTApi(self.properties, False)
        auth_provider = self.api_instance.rest_auth_function.auth_provider
        if isinstance(auth_provider, DLFAuthProvider):
            return auth_provider.get_token().access_key_id or 'anonymous'
        if isinstance(auth_provider, BearTokenAuthProvider):
            return auth_provider.token or 'anonymous'
        return 'anonymous'

    def try_to_refresh_token(self):
        identifier_str = self._build_cache_key()

        if self.token is not None and not self._is_token_expired(self.token):
            return

        cached_token = self._get_cached_token(identifier_str)
        if cached_token and not self._is_token_expired(cached_token):
            self.token = cached_token
            return

        global_lock = self._get_global_token_lock(identifier_str)

        with global_lock:
            cached_token = self._get_cached_token(identifier_str)
            if cached_token and not self._is_token_expired(cached_token):
                self.token = cached_token
                return

            token_to_check = cached_token if cached_token else self.token
            if token_to_check is None or self._is_token_expired(token_to_check):
                self.refresh_token()
                self._set_cached_token(identifier_str, self.token)

    def _get_cached_token(self, identifier_str: str) -> Optional[RESTToken]:
        with self._TOKEN_LOCKS_LOCK:
            return self._TOKEN_CACHE.get(identifier_str)

    def _set_cached_token(self, identifier_str: str, token: RESTToken):
        with self._TOKEN_LOCKS_LOCK:
            self._TOKEN_CACHE[identifier_str] = token

    def _is_token_expired(self, token: Optional[RESTToken]) -> bool:
        if token is None:
            return True
        current_time = int(time.time() * 1000)
        return (token.expire_at_millis - current_time) < RESTApi.TOKEN_EXPIRATION_SAFE_TIME_MILLIS

    def _get_global_token_lock(self, identifier_str: str) -> threading.Lock:
        with self._TOKEN_LOCKS_LOCK:
            if identifier_str not in self._TOKEN_LOCKS:
                self._TOKEN_LOCKS[identifier_str] = threading.Lock()
            return self._TOKEN_LOCKS[identifier_str]

    def refresh_token(self):
        self.log.info(f"begin refresh data token for identifier [{self.identifier}]")
        if self.api_instance is None:
            self.api_instance = RESTApi(self.properties, False)

        table_identifier = self.identifier
        if self.identifier.is_system_table():
            # Strip the system-table suffix; preserve the branch so the token
            # request resolves against the correct branch backing.
            table_identifier = Identifier.create(
                self.identifier.get_database_name(),
                self.identifier.get_table_name(),
                branch=self.identifier.get_branch_name(),
            )

        response = self.api_instance.load_table_token(table_identifier)
        self.log.info(
            f"end refresh data token for identifier [{self.identifier}] expiresAtMillis [{response.expires_at_millis}]"
        )

        merged_token_dict = self._merge_token_with_catalog_options(response.token)
        new_token = RESTToken(merged_token_dict, response.expires_at_millis)
        self.token = new_token

    def valid_token(self):
        self.try_to_refresh_token()
        return self.token

    def close(self):
        pass
