from __future__ import annotations

import uuid
from datetime import datetime
from enum import Enum, IntEnum
from typing import TYPE_CHECKING, Any, Dict, List, Mapping, Optional, Union

import numpy as np
from pydantic import BaseModel, Field, root_validator

from chalk.features.tag import EnvironmentId


def _category_for_error_code(c: Union[ErrorCode, str]) -> ErrorCodeCategory:
    c = ErrorCode[c]
    return {
        ErrorCode.PARSE_FAILED: ErrorCodeCategory.REQUEST,
        ErrorCode.RESOLVER_NOT_FOUND: ErrorCodeCategory.REQUEST,
        ErrorCode.INVALID_QUERY: ErrorCodeCategory.REQUEST,
        ErrorCode.VALIDATION_FAILED: ErrorCodeCategory.FIELD,
        ErrorCode.RESOLVER_FAILED: ErrorCodeCategory.FIELD,
        ErrorCode.RESOLVER_TIMED_OUT: ErrorCodeCategory.FIELD,
        ErrorCode.UPSTREAM_FAILED: ErrorCodeCategory.FIELD,
        ErrorCode.UNAUTHENTICATED: ErrorCodeCategory.NETWORK,
        ErrorCode.UNAUTHORIZED: ErrorCodeCategory.NETWORK,
        ErrorCode.INTERNAL_SERVER_ERROR: ErrorCodeCategory.NETWORK,
        ErrorCode.CANCELLED: ErrorCodeCategory.NETWORK,
        ErrorCode.DEADLINE_EXCEEDED: ErrorCodeCategory.NETWORK,
    }[c]


class OnlineQueryContext(BaseModel):
    """Context in which to execute a query."""

    environment: Optional[str] = None
    """
    The environment under which to run the resolvers.
    API tokens can be scoped to an # environment.
    If no environment is specified in the query,
    but the token supports only a single environment,
    then that environment will be taken as the scope
    for executing the request.
    """

    tags: Optional[List[str]] = None
    """
    The tags used to scope the resolvers.
    More information at https://docs.chalk.ai/docs/resolver-tags
    """


class OfflineQueryContext(BaseModel):
    environment: Optional[str] = None
    """
    The environment under which to run the resolvers.
    API tokens can be scoped to an # environment.
    If no environment is specified in the query,
    but the token supports only a single environment,
    then that environment will be taken as the scope
    for executing the request.
    """


class ErrorCode(str, Enum):
    """The detailed error code.

    For a simpler category of error, see `ErrorCodeCategory`.
    """

    PARSE_FAILED = "PARSE_FAILED"
    """The query contained features that do not exist."""

    RESOLVER_NOT_FOUND = "RESOLVER_NOT_FOUND"
    """
    A resolver was required as part of running the dependency
    graph that could not be found.
    """

    INVALID_QUERY = "INVALID_QUERY"
    """
    The query is invalid. All supplied features need to be
    rooted in the same top-level entity.
    """

    VALIDATION_FAILED = "VALIDATION_FAILED"
    """
    A feature value did not match the expected schema
    (e.g. `incompatible type "int"; expected "str"`)
    """

    RESOLVER_FAILED = "RESOLVER_FAILED"
    """The resolver for a feature errored."""

    RESOLVER_TIMED_OUT = "RESOLVER_TIMED_OUT"
    """The resolver for a feature timed out."""

    UPSTREAM_FAILED = "UPSTREAM_FAILED"
    """
    A crash in a resolver that was to produce an input for
    the resolver crashed, and so the resolver could not run
    crashed, and so the resolver could not run.
    """

    UNAUTHENTICATED = "UNAUTHENTICATED"
    """The request was submitted with an invalid authentication header."""

    UNAUTHORIZED = "UNAUTHORIZED"
    """The supplied credentials do not provide the right authorization to execute the request."""

    INTERNAL_SERVER_ERROR = "INTERNAL_SERVER_ERROR"
    """An unspecified error occurred."""

    CANCELLED = "CANCELLED"
    """The operation was cancelled, typically by the caller."""

    DEADLINE_EXCEEDED = "DEADLINE_EXCEEDED"
    """The deadline expired before the operation could complete."""


class ErrorCodeCategory(str, Enum):
    """The category of an error.

    For more detailed error information, see `ErrorCode`
    """

    REQUEST = "REQUEST"
    """
    Request errors are raised before execution of your
    resolver code. They may occur due to invalid feature
    names in the input or a request that cannot be satisfied
    by the resolvers you have defined.
    """

    FIELD = "FIELD"
    """
    Field errors are raised while running a feature resolver
    for a particular field. For this type of error, you'll
    find a feature and resolver attribute in the error type.
    When a feature resolver crashes, you will receive null
    value in the response. To differentiate from a resolver
    returning a null value and a failure in the resolver,
    you need to check the error schema.
    """

    NETWORK = "NETWORK"
    """
    Network errors are thrown outside your resolvers.
    For example, your request was unauthenticated,
    connection failed, or an error occurred within Chalk.
    """


class ChalkException(BaseModel, frozen=True):
    """Information about an exception from a resolver run."""

    kind: str
    """The name of the class of the exception."""

    message: str
    """The message taken from the exception."""

    stacktrace: str
    """The stacktrace produced by the code."""


class ChalkError(BaseModel, frozen=True):
    """
    The `ChalkError` describes an error from running a resolver
    or from a feature that can't be validated.
    """

    code: ErrorCode
    """The type of the error."""

    category: ErrorCodeCategory = ErrorCodeCategory.NETWORK
    """
    The category of the error, given in the type field for the error codes.
    This will be one of "REQUEST", "NETWORK", and "FIELD".
    """

    message: str
    """A readable description of the error message."""

    exception: Optional[ChalkException] = None
    """The exception that caused the failure, if applicable."""

    feature: Optional[str] = None
    """
    The fully qualified name of the failing feature, e.g. `user.identity.has_voip_phone`.
    """

    resolver: Optional[str] = None
    """
    The fully qualified name of the failing resolver, e.g. `my.project.get_fraud_score`.
    """

    def copy_for_feature(self, feature: str) -> "ChalkError":
        return self.copy(update={"feature": feature})

    @root_validator
    def _validate_category(cls, values: Dict[str, Any]):
        values["category"] = _category_for_error_code(values["code"])
        return values

    if TYPE_CHECKING:
        # Defining __hash__ only when type checking
        # since pydantic provides a hash for frozen models
        def __hash__(self) -> int:  # type: ignore
            ...


class ResolverRunStatus(str, Enum):
    """Status of a scheduled resolver run."""

    RECEIVED = "received"
    """The request to run the resolver has been received, and is running or scheduled."""

    SUCCEEDED = "succeeded"
    """The resolver run failed."""

    FAILED = "failed"
    """The resolver run succeeded."""


class ResolverRunResponse(BaseModel):
    """Status of a scheduled resolver run."""

    id: str
    """The ID of the resolver run."""

    status: ResolverRunStatus
    """The current status of the resolver run."""


class WhoAmIResponse(BaseModel):
    """Response for checking the authenticated user."""

    user: str
    """The ID of the user or service token making the query."""


class FeatureResolutionMeta(BaseModel, frozen=True):
    """Detailed metadata about the execution of an online query."""

    chosen_resolver_fqn: str
    """The name of the resolver that computed the feature value."""

    cache_hit: bool
    """Whether the feature request was satisfied by a cached value."""

    primitive_type: Optional[str] = None
    """
    Primitive type name for the feature, e.g. `str` for `some_feature: str`.
    Returned only if query-level 'include_meta' is True.
    """

    version: int = 1
    """
    The version that was selected for this feature. Defaults to `default_version`, if query
    does not specify a constraint. If no versioning information is provided on the feature definition,
    the default version is `1`.
    """


class FeatureResult(BaseModel):
    field: str
    """
    The name of the feature requested, e.g. 'user.identity.has_voip_phone'.
    """

    value: Any  # Value should be a TJSON type
    """
    The value of the requested feature.
    If an error was encountered in resolving this feature,
    this field will be empty.
    """

    pkey: Optional[Union[str, int]]
    """The primary key of the resolved feature."""

    error: Optional[ChalkError] = None
    """
    The error code encountered in resolving this feature.
    If no error occurred, this field is empty.
    """

    ts: Optional[datetime] = None
    """
    The time at which this feature was computed.
    This value could be significantly in the past if you're using caching.
    """

    meta: Optional[FeatureResolutionMeta] = None
    """Detailed information about how this feature was computed."""


class ExchangeCredentialsRequest(BaseModel):
    client_id: str
    client_secret: str
    grant_type: str


class ExchangeCredentialsResponse(BaseModel):
    access_token: str
    token_type: str
    expires_in: int
    # expires_at: datetime
    api_server: str
    primary_environment: Optional[str] = None
    engines: Optional[Mapping[str, str]] = None


class OfflineQueryInput(BaseModel):
    columns: List[str]
    values: List[List[Any]]  # Values should be of type TJSON


class OnlineQueryRequest(BaseModel):
    inputs: Mapping[str, Any]  # Values should be of type TJSON
    outputs: List[str]
    staleness: Optional[Mapping[str, str]] = None
    context: Optional[OnlineQueryContext] = None
    include_meta: bool = True
    explain: bool = False
    skip_online_storage: Optional[bool] = False
    skip_offline_storage: Optional[bool] = False
    skip_metrics_storage: Optional[bool] = False
    skip_cache_lookups: Optional[bool] = False
    correlation_id: Optional[str] = None
    query_name: Optional[str] = None
    deployment_id: Optional[str] = None
    branch_id: Optional[str] = None
    meta: Optional[Mapping[str, str]] = None


class TriggerResolverRunRequest(BaseModel):
    resolver_fqn: str


class QueryMeta(BaseModel):
    execution_duration_s: float
    """
    The time, expressed in seconds, that Chalk spent executing this query.
    """

    deployment_id: Optional[str] = None
    """
    The id of the deployment that served this query.
    """

    environment_id: Optional[str] = None
    """
    The id of the environment that served this query. Not intended to be human readable, but helpful for support.
    """

    environment_name: Optional[str] = None
    """
    The short name of the environment that served this query. For example: "dev" or "prod".
    """

    query_id: Optional[str] = None
    """
    A unique ID generated and persisted by Chalk for this query. All computed features, metrics, and logs are
    associated with this ID. Your system can store this ID for audit and debugging workflows.
    """

    query_timestamp: Optional[datetime] = None
    """
    At the start of query execution, Chalk computes 'datetime.now()'. This value is used to timestamp computed features.
    """

    query_hash: Optional[str] = None
    """
    Deterministic hash of the 'structure' of the query. Queries that have the same input/output features will
    typically have the same hash; changes may be observed over time as we adjust implementation details.
    """

    explain_output: Optional[str] = None
    """
    An unstructured string containing diagnostic information about the query execution. Only included if `explain` is True.
    """


class OnlineQueryResponse(BaseModel):
    data: List[FeatureResult]
    errors: Optional[List[ChalkError]] = None
    meta: Optional[QueryMeta] = None

    def for_fqn(self, fqn: str):
        return next((x for x in self.data if x.field == fqn), None)

    class Config:
        json_encoders = {
            np.integer: int,
            np.floating: float,
        }


class CreateOfflineQueryJobRequest(BaseModel):
    output: List[str] = Field(description="A list of output feature root fqns to query")
    required_output: List[str] = Field(default_factory=list, description="A list of required output feature root fqns")
    destination_format: str = Field(description="The desired output format. Should be 'CSV' or 'PARQUET'")
    job_id: Optional[uuid.UUID] = Field(
        default=None,
        description=(
            "A unique job id. If not specified, one will be auto generated by the server. If specified by the client, "
            "then jobs with the same ID will be rejected."
        ),
    )
    input: Optional[OfflineQueryInput] = Field(default=None, description="Any givens")
    max_samples: Optional[int] = Field(
        default=None,
        description="The maximum number of samples. If None, no limit",
    )
    max_cache_age_secs: Optional[int] = Field(
        default=None,  # Defaults to ``OFFLINE_QUERY_MAX_CACHE_AGE_SECS`` in the chalkengine config
        description=(
            "The maximum staleness, in seconds, for how old the view on the offline store can be. That is, "
            "data ingested within this interval will not be reflected in this offline query. "
            "Set to ``0`` to ignore the cache. If not specified, it defaults to 30 minutes."
        ),
    )
    observed_at_lower_bound: Optional[datetime] = Field(
        default=None,
        description="The lower bound for the observed at timestamp (inclusive). If not specified, defaults to the beginning of time",
    )
    observed_at_upper_bound: Optional[datetime] = Field(
        default=None,
        description="The upper bound for the observed at timestamp (inclusive). If not specified, defaults to the end of time.",
    )
    dataset_name: Optional[str] = None
    branch: Optional[str] = None


class ComputeResolverOutputRequest(BaseModel):
    input: OfflineQueryInput
    resolver_fqn: str
    branch: Optional[str] = None
    environment: Optional[str] = None


class DatasetRecomputeRequest(BaseModel):
    dataset_name: Optional[str] = None
    branch: str
    dataset_id: Optional[str]
    revision_id: Optional[str]
    features: List[str]


class RecomputeResolverOutputRequest(BaseModel):
    persistent_id: str
    resolver_fqn: str
    branch: Optional[str] = None
    environment: Optional[str] = None


class ComputeResolverOutputResponse(BaseModel):
    job_id: str
    persistent_id: str
    errors: Optional[List[ChalkError]] = None


class OfflineQueryRequest(BaseModel):
    """V1 OfflineQueryRequest. Not used by the current Chalk Client."""

    output: List[str]  # output features which can be null
    input: Optional[OfflineQueryInput] = None
    dataset: Optional[str] = None
    max_samples: Optional[int] = None
    max_cache_age_secs: Optional[int] = None
    required_outputs: List[str] = Field(default_factory=list)  # output features which cannot be null


class OfflineQueryResponse(BaseModel):
    """V1 OfflineQueryResponse. Not used by the current Chalk Client."""

    columns: List[str]
    output: List[List[Any]]  # values should be of TJSON types
    errors: Optional[List[ChalkError]] = None


class CreateOfflineQueryJobResponse(BaseModel):
    """
    Attributes:
        job_id: A job ID, which can be used to retrieve the results.
    """

    job_id: uuid.UUID
    version: int = 1  # Field is deprecated
    errors: Optional[List[ChalkError]] = None


class ColumnMetadata(BaseModel):
    feature_fqn: str = Field(description="The root FQN of the feature for a column")

    column_name: str = Field(description="The name of the column that corresponds to this feature")

    dtype: str = Field(description="The data type for this feature")
    # This field is currently a JSON-stringified version of the SerializeDType property
    # Using a string instead of a pydantic model the SerializedDType encoding does not affect
    # the api layer


class GetOfflineQueryJobResponse(BaseModel):
    is_finished: bool = Field(description="Whether the export job is finished (it runs asynchronously)")
    version: int = Field(
        default=1,  # Backwards compatibility
        description=(
            "Version number representing the format of the data. The client uses this version number "
            "to properly decode and load the query results into DataFrames."
        ),
    )
    urls: List[str] = Field(
        description="A list of short-lived, authenticated URLs that the client can download to retrieve the exported data."
    )
    errors: Optional[List[ChalkError]] = None
    columns: Optional[List[ColumnMetadata]] = Field(
        description="Expected columns for the dataframe, including data type information",
        default=None,
    )


class QueryStatus(IntEnum):
    PENDING_SUBMISSION = 1
    """Pending submission to the database."""
    SUBMITTED = 2
    """Submitted to the database, but not yet running."""
    RUNNING = 3
    """Running in the database."""
    ERROR = 4
    """Error with either submitting or running the job."""
    EXPIRED = 5
    """The job did not complete before an expiration deadline, so there are no results."""
    CANCELLED = 6
    """Manually cancelled before it errored or finished successfully."""
    SUCCESSFUL = 7  #
    """Successfully ran the job."""


class DatasetSampleFilter(BaseModel):
    lower_bound: Optional[datetime] = None
    upper_bound: Optional[datetime] = None
    max_samples: Optional[int] = None


class DatasetFilter(BaseModel):
    sample_filters: DatasetSampleFilter = Field(default_factory=DatasetSampleFilter)
    max_cache_age_secs: Optional[float] = None


class DatasetRevisionResponse(BaseModel):
    dataset_name: Optional[str] = None
    dataset_id: Optional[uuid.UUID] = None
    environment_id: EnvironmentId
    revision_id: uuid.UUID  # Currently, the revision ID is the job ID that created the revision
    creator_id: str
    outputs: List[str]
    created_at: Optional[datetime] = None
    started_at: Optional[datetime] = None
    terminated_at: Optional[datetime] = None
    givens_uri: Optional[str] = None
    status: QueryStatus
    filters: DatasetFilter
    num_partitions: int
    num_bytes: Optional[int] = None
    output_uris: str
    output_version: int
    branch: Optional[str] = None


class DatasetResponse(BaseModel):
    is_finished: bool = Field(description="Whether the export job is finished (it runs asynchronously)")
    version: int = Field(
        default=1,  # Backwards compatibility
        description=(
            "Version number representing the format of the data. The client uses this version number "
            "to properly decode and load the query results into DataFrames."
        ),
    )
    environment_id: EnvironmentId
    dataset_id: Optional[uuid.UUID] = None
    dataset_name: Optional[str] = None
    revisions: List[DatasetRevisionResponse]
    errors: Optional[List[ChalkError]] = None


class UpdateResolverResponse(BaseModel):
    updated_fqn: Optional[
        str
    ] = None  # The resolver fqn that was updated (may not be the same as the one that was requested)
    is_new: Optional[bool] = None  # Whether a new resolver was created, or if an existing one was replaced
    errors: Optional[List[ChalkError]] = None
