import logging
import os
import warnings
from pathlib import Path
from typing import Any, Dict, Optional

import yaml
from pydantic import (
    BaseModel,
    ConfigDict,
    Field,
    StrictBool,
    StrictInt,
    StrictStr,
    ValidationError,
    ValidationInfo,
    field_validator,
    model_validator,
)

from feast.errors import (
    FeastFeatureServerTypeInvalidError,
    FeastInvalidAuthConfigClass,
    FeastOfflineStoreInvalidName,
    FeastOnlineStoreInvalidName,
    FeastRegistryNotSetError,
    FeastRegistryTypeInvalidError,
)
from feast.importer import import_class
from feast.permissions.auth.auth_type import AuthType

warnings.simplefilter("once", RuntimeWarning)

_logger = logging.getLogger(__name__)

# These dict exists so that:
# - existing values for the online store type in featurestore.yaml files continue to work in a backwards compatible way
# - first party and third party implementations can use the same class loading code path.
REGISTRY_CLASS_FOR_TYPE = {
    "file": "feast.infra.registry.registry.Registry",
    "sql": "feast.infra.registry.sql.SqlRegistry",
    "snowflake.registry": "feast.infra.registry.snowflake.SnowflakeRegistry",
    "remote": "feast.infra.registry.remote.RemoteRegistry",
}

BATCH_ENGINE_CLASS_FOR_TYPE = {
    "local": "feast.infra.materialization.local_engine.LocalMaterializationEngine",
    "snowflake.engine": "feast.infra.materialization.snowflake_engine.SnowflakeMaterializationEngine",
    "lambda": "feast.infra.materialization.aws_lambda.lambda_engine.LambdaMaterializationEngine",
    "k8s": "feast.infra.materialization.kubernetes.k8s_materialization_engine.KubernetesMaterializationEngine",
    "spark.engine": "feast.infra.materialization.contrib.spark.spark_materialization_engine.SparkMaterializationEngine",
}

LEGACY_ONLINE_STORE_CLASS_FOR_TYPE = {
    "feast.infra.online_stores.contrib.postgres.PostgreSQLOnlineStore": "feast.infra.online_stores.postgres_online_store.PostgreSQLOnlineStore",
    "feast.infra.online_stores.contrib.hbase_online_store.hbase.HbaseOnlineStore": "feast.infra.online_stores.hbase_online_store.hbase.HbaseOnlineStore",
    "feast.infra.online_stores.contrib.cassandra_online_store.cassandra_online_store.CassandraOnlineStore": "feast.infra.online_stores.cassandra_online_store.cassandra_online_store.CassandraOnlineStore",
    "feast.infra.online_stores.contrib.mysql_online_store.mysql.MySQLOnlineStore": "feast.infra.online_stores.mysql_online_store.mysql.MySQLOnlineStore",
    "feast.infra.online_stores.contrib.hazelcast_online_store.hazelcast_online_store.HazelcastOnlineStore": "feast.infra.online_stores.hazelcast_online_store.hazelcast_online_store.HazelcastOnlineStore",
    "feast.infra.online_stores.contrib.ikv_online_store.ikv.IKVOnlineStore": "feast.infra.online_stores.ikv_online_store.ikv.IKVOnlineStore",
    "feast.infra.online_stores.contrib.elasticsearch.ElasticSearchOnlineStore": "feast.infra.online_stores.elasticsearch_online_store.ElasticSearchOnlineStore",
    "feast.infra.online_stores.contrib.singlestore_online_store.singlestore.SingleStoreOnlineStore": "feast.infra.online_stores.singlestore_online_store.singlestore.SingleStoreOnlineStore",
    "feast.infra.online_stores.contrib.qdrant.QdrantOnlineStore": "feast.infra.online_stores.cqdrant.QdrantOnlineStore",
}

ONLINE_STORE_CLASS_FOR_TYPE = {
    "sqlite": "feast.infra.online_stores.sqlite.SqliteOnlineStore",
    "datastore": "feast.infra.online_stores.datastore.DatastoreOnlineStore",
    "redis": "feast.infra.online_stores.redis.RedisOnlineStore",
    "dynamodb": "feast.infra.online_stores.dynamodb.DynamoDBOnlineStore",
    "snowflake.online": "feast.infra.online_stores.snowflake.SnowflakeOnlineStore",
    "bigtable": "feast.infra.online_stores.bigtable.BigtableOnlineStore",
    "postgres": "feast.infra.online_stores.postgres_online_store.postgres.PostgreSQLOnlineStore",
    "hbase": "feast.infra.online_stores.hbase_online_store.hbase.HbaseOnlineStore",
    "cassandra": "feast.infra.online_stores.cassandra_online_store.cassandra_online_store.CassandraOnlineStore",
    "mysql": "feast.infra.online_stores.mysql_online_store.mysql.MySQLOnlineStore",
    "hazelcast": "feast.infra.online_stores.hazelcast_online_store.hazelcast_online_store.HazelcastOnlineStore",
    "ikv": "feast.infra.online_stores.ikv_online_store.ikv.IKVOnlineStore",
    "elasticsearch": "feast.infra.online_stores.elasticsearch_online_store.ElasticSearchOnlineStore",
    "remote": "feast.infra.online_stores.remote.RemoteOnlineStore",
    "singlestore": "feast.infra.online_stores.singlestore_online_store.singlestore.SingleStoreOnlineStore",
    "qdrant": "feast.infra.online_stores.cqdrant.QdrantOnlineStore",
    "couchbase": "feast.infra.online_stores.couchbase_online_store.couchbase.CouchbaseOnlineStore",
    "milvus": "feast.infra.online_stores.milvus_online_store.milvus.MilvusOnlineStore",
    **LEGACY_ONLINE_STORE_CLASS_FOR_TYPE,
}

OFFLINE_STORE_CLASS_FOR_TYPE = {
    "file": "feast.infra.offline_stores.dask.DaskOfflineStore",
    "dask": "feast.infra.offline_stores.dask.DaskOfflineStore",
    "bigquery": "feast.infra.offline_stores.bigquery.BigQueryOfflineStore",
    "redshift": "feast.infra.offline_stores.redshift.RedshiftOfflineStore",
    "snowflake.offline": "feast.infra.offline_stores.snowflake.SnowflakeOfflineStore",
    "spark": "feast.infra.offline_stores.contrib.spark_offline_store.spark.SparkOfflineStore",
    "trino": "feast.infra.offline_stores.contrib.trino_offline_store.trino.TrinoOfflineStore",
    "postgres": "feast.infra.offline_stores.contrib.postgres_offline_store.postgres.PostgreSQLOfflineStore",
    "athena": "feast.infra.offline_stores.contrib.athena_offline_store.athena.AthenaOfflineStore",
    "mssql": "feast.infra.offline_stores.contrib.mssql_offline_store.mssql.MsSqlServerOfflineStore",
    "duckdb": "feast.infra.offline_stores.duckdb.DuckDBOfflineStore",
    "remote": "feast.infra.offline_stores.remote.RemoteOfflineStore",
}

FEATURE_SERVER_CONFIG_CLASS_FOR_TYPE = {
    "local": "feast.infra.feature_servers.local_process.config.LocalFeatureServerConfig",
}

ALLOWED_AUTH_TYPES = ["no_auth", "kubernetes", "oidc"]

AUTH_CONFIGS_CLASS_FOR_TYPE = {
    "no_auth": "feast.permissions.auth_model.NoAuthConfig",
    "kubernetes": "feast.permissions.auth_model.KubernetesAuthConfig",
    "oidc": "feast.permissions.auth_model.OidcAuthConfig",
    "oidc_client": "feast.permissions.auth_model.OidcClientAuthConfig",
}


class FeastBaseModel(BaseModel):
    """Feast Pydantic Configuration Class"""

    model_config = ConfigDict(arbitrary_types_allowed=True, extra="allow")


class FeastConfigBaseModel(BaseModel):
    """Feast Pydantic Configuration Class"""

    model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid")


class RegistryConfig(FeastBaseModel):
    """Metadata Store Configuration. Configuration that relates to reading from and writing to the Feast registry."""

    registry_type: StrictStr = "file"
    """ str: Provider name or a class name that implements Registry."""

    registry_store_type: Optional[StrictStr] = None
    """ str: Provider name or a class name that implements RegistryStore. """

    path: StrictStr = ""
    """ str: Path to metadata store.
        If registry_type is 'file', then an be a local path, or remote object storage path, e.g. a GCS URI
        If registry_type is 'sql', then this is a database URL as expected by SQLAlchemy """

    cache_ttl_seconds: StrictInt = 600
    """int: The cache TTL is the amount of time registry state will be cached in memory. If this TTL is exceeded then
     the registry will be refreshed when any feature store method asks for access to registry state. The TTL can be
     set to infinity by setting TTL to 0 seconds, which means the cache will only be loaded once and will never
     expire. Users can manually refresh the cache by calling feature_store.refresh_registry() """

    s3_additional_kwargs: Optional[Dict[str, str]] = None
    """ Dict[str, str]: Extra arguments to pass to boto3 when writing the registry file to S3. """

    purge_feast_metadata: StrictBool = False
    """ bool: Stops using feast_metadata table and delete data from feast_metadata table.
        Once this is set to True, it cannot be reverted back to False. Reverting back to False will
        only reset the project but not all the projects"""

    @field_validator("path")
    def validate_path(cls, path: str, values: ValidationInfo) -> str:
        if values.data.get("registry_type") == "sql":
            if path.startswith("postgresql://"):
                _logger.warning(
                    "The `path` of the `RegistryConfig` starts with a plain "
                    "`postgresql` string. We are updating this to `postgresql+psycopg` "
                    "to ensure that the `psycopg3` driver is used by `sqlalchemy`. If "
                    "you want to use `psycopg2` pass `postgresql+psycopg2` explicitely "
                    "to `path`. To silence this warning, pass `postgresql+psycopg` "
                    "explicitely to `path`."
                )
                return path.replace("postgresql://", "postgresql+psycopg://")
        return path


class RepoConfig(FeastBaseModel):
    """Repo config. Typically loaded from `feature_store.yaml`"""

    project: StrictStr
    """ str: This acts as a Feast unique project identifier. This can be any alphanumeric string and can have '_' but can not start with '_'.
        You can have multiple independent feature repositories deployed to the same cloud
        provider account, as long as they have different project identifier.
    """

    provider: StrictStr = "local"
    """ str: local or gcp or aws """

    registry_config: Any = Field(alias="registry", default="data/registry.db")
    """ Configures the registry.
        Can be:
            1. str: a path to a file based registry (a local path, or remote object storage path, e.g. a GCS URI)
            2. RegistryConfig: A fully specified file based registry or SQL based registry
            3. SnowflakeRegistryConfig: Using a Snowflake table to store the registry
    """

    online_config: Any = Field(None, alias="online_store")
    """ OnlineStoreConfig: Online store configuration (optional depending on provider) """

    auth: Any = Field(None, alias="auth")
    """ auth: Optional if the services needs the authentication against IDPs (optional depending on provider) """

    offline_config: Any = Field(None, alias="offline_store")
    """ OfflineStoreConfig: Offline store configuration (optional depending on provider) """

    batch_engine_config: Any = Field(None, alias="batch_engine")
    """ BatchMaterializationEngine: Batch materialization configuration (optional depending on provider)"""

    feature_server: Optional[Any] = None
    """ FeatureServerConfig: Feature server configuration (optional depending on provider) """

    flags: Any = None
    """ Flags (deprecated field): Feature flags for experimental features """

    repo_path: Optional[Path] = None
    """When using relative path in FileSource path, this parameter is mandatory"""

    entity_key_serialization_version: StrictInt = 1
    """ Entity key serialization version: This version is used to control what serialization scheme is
    used when writing data to the online store.
    A value