firn 0.15.0

use std::collections::HashMap;

use serde::Deserialize;

#[allow(clippy::large_enum_variant)]
#[derive(Deserialize, Debug)]
#[serde(untagged)]
pub enum ExecResponse {
    // QueryAsync must come before Query: a 333334 in-progress body only contains
    // {queryId, getResultUrl, ...}, which is a strict subset of the fields
    // QueryExecResponseData allows via Option<>, so untagged would otherwise
    // greedily match the wrong arm.
    QueryAsync(QueryAsyncExecResponse),
    Query(QueryExecResponse),
    PutGet(PutGetExecResponse),
    Error(ExecErrorResponse),
}

/// True for both `333333` (queryInProgressCode) and `333334`
/// (queryInProgressAsyncCode). gosnowflake distinguishes them; both indicate
/// the result is not ready and we should keep polling.
pub fn is_query_in_progress(code: Option<&String>) -> bool {
    matches!(code.map(String::as_str), Some("333333" | "333334"))
}

/// `390112` — gosnowflake's `sessionExpiredCode`. Returned mid-flight when
/// the session token has expired; the caller should renew and retry.
pub fn is_session_expired(code: Option<&String>) -> bool {
    matches!(code.map(String::as_str), Some("390112"))
}

/// `000605` — gosnowflake's `queryNotExecutingCode`. Returned by the abort
/// endpoint when there is no in-flight query to cancel (for example, the
/// query already finished). Treat as success.
pub fn is_query_not_executing(code: Option<&String>) -> bool {
    matches!(code.map(String::as_str), Some("000605"))
}

/// `000604` — Snowflake's "SQL execution canceled" code. The original
/// query's polling/initial-POST request observes this *after* an abort
/// has been sent (typically from a different task or endpoint), so we
/// can map it to a clean `QueryCancelled` instead of a generic API error.
pub fn is_sql_execution_cancelled(code: Option<&String>) -> bool {
    matches!(code.map(String::as_str), Some("000604"))
}

// todo: add close session response, which should be just empty?
// FIXME: dead_code
#[allow(clippy::large_enum_variant, dead_code)]
#[derive(Deserialize, Debug)]
#[serde(untagged)]
pub enum AuthResponse {
    Login(LoginResponse),
    Auth(AuthenticatorResponse),
    Renew(RenewSessionResponse),
    Close(CloseSessionResponse),
    Error(AuthErrorResponse),
}

#[derive(Deserialize, Debug)]
pub struct BaseRestResponse<D> {
    // null for auth
    pub code: Option<String>,
    pub message: Option<String>,
    pub success: bool,
    pub data: D,
}

pub type PutGetExecResponse = BaseRestResponse<PutGetResponseData>;
pub type QueryExecResponse = BaseRestResponse<QueryExecResponseData>;
pub type QueryAsyncExecResponse = BaseRestResponse<QueryAsyncExecResponseData>;
pub type ExecErrorResponse = BaseRestResponse<ExecErrorResponseData>;
/// Response from `/queries/v1/abort-request`. Snowflake returns `data: null`
/// (or sometimes a sparse object); we keep it as a free-form Value so the
/// parser doesn't choke.
pub type CancelQueryResponse = BaseRestResponse<serde_json::Value>;
/// Response from `GET /monitoring/queries/{queryId}`. Used by
/// `SnowflakeApi::query_status` to peek at a query without consuming
/// warehouse credits or buffering results.
pub type MonitoringResponse = BaseRestResponse<MonitoringResponseData>;
pub type AuthErrorResponse = BaseRestResponse<AuthErrorResponseData>;
pub type AuthenticatorResponse = BaseRestResponse<AuthenticatorResponseData>;
pub type LoginResponse = BaseRestResponse<LoginResponseData>;
pub type RenewSessionResponse = BaseRestResponse<RenewSessionResponseData>;
// Data should be always `null` on successful close session response
pub type CloseSessionResponse = BaseRestResponse<Option<()>>;

#[derive(Deserialize, Debug)]
#[serde(rename_all = "camelCase")]
pub struct ExecErrorResponseData {
    pub age: i64,
    pub error_code: String,
    // Optional because the post-abort error response (code 000604 "SQL
    // execution canceled") omits this field.
    pub internal_error: Option<bool>,

    // come when query is invalid
    pub line: Option<i64>,
    pub pos: Option<i64>,

    // Optional because the abort-request error response and some session
    // errors omit these. Per spiceai-1.1.
    pub query_id: Option<String>,
    pub sql_state: Option<String>,
}

#[derive(Deserialize, Debug)]
#[serde(rename_all = "camelCase")]
pub struct MonitoringResponseData {
    /// In practice Snowflake returns either zero (unknown id) or one entry
    /// for the requested `query_id`. The list shape mirrors the wire format.
    #[serde(default)]
    pub queries: Vec<MonitoringQueryEntry>,
}

#[derive(Deserialize, Debug)]
#[serde(rename_all = "camelCase")]
#[allow(dead_code)]
pub struct MonitoringQueryEntry {
    /// Echoed `queryId`.
    pub id: String,
    /// gosnowflake's wire string, e.g. `"RUNNING"`, `"SUCCESS"`,
    /// `"FAILED_WITH_ERROR"`. Always present per gosnowflake.
    pub status: String,
    pub error_code: Option<String>,
    pub error_message: Option<String>,
}

#[derive(Deserialize, Debug)]
#[serde(rename_all = "camelCase")]
pub struct QueryAsyncExecResponseData {
    pub query_id: String,
    pub get_result_url: String,
    pub query_aborts_after_secs: Option<i64>,
    pub progress_desc: Option<String>,
}

#[derive(Deserialize, Debug)]
#[serde(rename_all = "camelCase")]
// FIXME: dead_code
#[allow(dead_code)]
pub struct AuthErrorResponseData {
    pub authn_method: Option<String>,
    pub error_code: Option<String>,
}

#[derive(Deserialize, Debug)]
pub struct NameValueParameter {
    pub name: String,
    pub value: serde_json::Value,
}

#[derive(Deserialize, Debug)]
#[serde(rename_all = "camelCase")]
// FIXME
#[allow(dead_code)]
pub struct LoginResponseData {
    pub session_id: i64,
    pub token: String,
    pub master_token: String,
    pub server_version: String,
    #[serde(default)]
    pub parameters: Vec<NameValueParameter>,
    pub session_info: SessionInfo,
    pub master_validity_in_seconds: i64,
    pub validity_in_seconds: i64,
}

#[derive(Deserialize, Debug)]
#[serde(rename_all = "camelCase")]
// FIXME: dead_code
#[allow(dead_code)]
pub struct SessionInfo {
    pub database_name: Option<String>,
    pub schema_name: Option<String>,
    pub warehouse_name: Option<String>,
    pub role_name: String,
}

#[derive(Deserialize, Debug)]
#[serde(rename_all = "camelCase")]
// FIXME: dead_code
#[allow(dead_code)]
pub struct AuthenticatorResponseData {
    pub token_url: Option<String>,
    pub sso_url: String,
    pub proof_key: String,
}

#[derive(Deserialize, Debug)]
#[serde(rename_all = "camelCase")]
// FIXME: dead_code
#[allow(dead_code)]
pub struct RenewSessionResponseData {
    pub session_token: String,
    pub validity_in_seconds_s_t: i64,
    pub master_token: String,
    pub validity_in_seconds_m_t: i64,
    pub session_id: i64,
}

#[derive(Deserialize, Debug)]
#[serde(rename_all = "camelCase")]
pub struct QueryExecResponseData {
    pub parameters: Vec<NameValueParameter>,
    pub rowtype: Vec<ExecResponseRowType>,
    // default for non-SELECT queries
    // GET / PUT has their own response format
    pub rowset: Option<serde_json::Value>,
    // only exists when binary response is given, eg Arrow
    // default for all SELECT queries
    // is base64-encoded Arrow IPC payload
    pub rowset_base64: Option<String>,
    pub total: i64,
    pub returned: i64,    // unused in .NET
    pub query_id: String, // unused in .NET
    pub database_provider: Option<String>,
    pub final_database_name: Option<String>, // unused in .NET
    pub final_schema_name: Option<String>,
    pub final_warehouse_name: Option<String>, // unused in .NET
    pub final_role_name: String,              // unused in .NET
    // only present on SELECT queries
    pub number_of_binds: Option<i32>, // unused in .NET
    // todo: deserialize into enum
    pub statement_type_id: i64,
    pub version: i64,
    // if response is chunked
    #[serde(default)] // soft-default to empty Vec if not present
    pub chunks: Vec<ExecResponseChunk>,
    // x-amz-server-side-encryption-customer-key, when chunks are present for download
    pub qrmk: Option<String>,
    #[serde(default)] // chunks are present
    pub chunk_headers: HashMap<String, String>,
    // when async query is run (ping pong request?)
    pub get_result_url: Option<String>,
    // multi-statement response, comma-separated
    pub result_ids: Option<String>,
    // `progressDesc`, and `queryAbortAfterSecs` are not used but exist in .NET
    // `sendResultTime`, `queryResultFormat`, `queryContext` also exist
}

#[derive(Deserialize, Debug, Clone)]
#[serde(rename_all = "camelCase")]
pub struct ExecResponseRowType {
    // Top-level rowtype entries always carry a name. Nested entries
    // (VECTOR's element-type sub-schema) come back without one — let
    // serde fall back to "" rather than fail the whole deserialize.
    #[serde(default)]
    pub name: String,
    pub byte_length: Option<i64>,
    // unused in .NET
    pub length: Option<i64>,
    #[serde(rename = "type")]
    pub type_: SnowflakeType,
    pub scale: Option<i64>,
    pub precision: Option<i64>,
    pub nullable: bool,
    /// Snowflake's extension-type discriminator. For columns whose
    /// canonical wire `type` is a generic catch-all, the *real* type
    /// shows up here as an upper-case tag. Observed values include
    /// `"GEOGRAPHY"` and `"GEOMETRY"` (both ride on `type: "object"`);
    /// callers that care about the difference should match on this
    /// rather than `type_`. None for plain types.
    #[serde(default)]
    pub ext_type_name: Option<String>,
    /// Dimensionality for `VECTOR(<element>, N)` columns. Carried in a
    /// dedicated field rather than `length`/`precision`. None for
    /// non-vector types.
    #[serde(default)]
    pub vector_dimension: Option<i64>,
    /// Nested sub-schema for parametric types. For `VECTOR(FLOAT, 3)`
    /// this carries one entry describing the element type
    /// (`type: "real"` etc.). None / empty for non-parametric types.
    #[serde(default)]
    pub fields: Option<Vec<ExecResponseRowType>>,
}

// fixme: is it good idea to keep this as an enum if more types could be added in future?
#[derive(Deserialize, Debug, Clone, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
pub enum SnowflakeType {
    Fixed,
    Real,
    Text,
    Date,
    Variant,
    TimestampLtz,
    TimestampNtz,
    TimestampTz,
    Object,
    Binary,
    Time,
    Boolean,
    Array,
    /// `VECTOR(<element>, N)`. Wire tag is plain `"vector"`; the
    /// dimension and element type live on
    /// `ExecResponseRowType::vector_dimension` and `fields[0]`.
    /// `GEOGRAPHY` and `GEOMETRY` ride on [`Self::Object`] with the
    /// real type in `ext_type_name` — they don't need their own
    /// variants here.
    Vector,
    /// `MAP(<k>, <v>)`. Listed in gosnowflake's converter type table
    /// (`MapType`) but empirically Snowflake collapses MAP results to
    /// plain `"object"` on the wire today — this variant exists so
    /// deserialization doesn't blow up if the protocol ever surfaces a
    /// distinct `"map"` tag. Defensive parity with gosnowflake.
    Map,
    /// IEEE-754 decimal (`DECFLOAT`). Listed in gosnowflake's converter
    /// table (`DecfloatType`) but at result time Snowflake currently
    /// demotes DECFLOAT casts to `"text"` on the wire. Defensive
    /// variant in case the protocol ever surfaces a distinct
    /// `"decfloat"` tag. Parity with gosnowflake.
    Decfloat,
}

#[derive(Deserialize, Debug)]
#[serde(rename_all = "camelCase")]
pub struct ExecResponseChunk {
    pub url: String,
    pub row_count: i32,
    pub uncompressed_size: i64,
}

#[derive(Deserialize, Debug)]
#[serde(rename_all = "camelCase")]
pub struct PutGetResponseData {
    pub query_id: String,
    // `kind`, `operation` are present in Go implementation, but not in .NET
    pub command: CommandType,
    pub local_location: Option<String>,
    // inconsistent case naming
    #[serde(rename = "src_locations", default)]
    pub src_locations: Vec<String>,
    // file upload parallelism
    pub parallel: usize, // fixme: originally i32, handle this in parsing somehow?
    // file size threshold, small ones are should be uploaded with given parallelism
    pub threshold: i64,
    // doesn't need compression if source is already compressed
    pub auto_compress: bool,
    pub overwrite: bool,
    // maps to one of the predefined compression algos
    // todo: support different compression formats?
    pub source_compression: String,
    pub stage_info: PutGetStageInfo,
    pub encryption_material: EncryptionMaterialVariant,
    // GCS specific. If you request multiple files?
    #[serde(default)]
    pub presigned_urls: Vec<String>,
    #[serde(default)]
    pub parameters: Vec<NameValueParameter>,
    pub statement_type_id: Option<i64>,
}

#[derive(Deserialize, Debug)]
#[serde(rename_all = "UPPERCASE")]
pub enum CommandType {
    Upload,
    Download,
}

#[derive(Deserialize, Debug)]
#[serde(untagged)]
pub enum PutGetStageInfo {
    Aws(AwsPutGetStageInfo),
    Azure(AzurePutGetStageInfo),
    Gcs(GcsPutGetStageInfo),
}

#[derive(Deserialize, Debug)]
#[serde(rename_all = "camelCase")]
pub struct AwsPutGetStageInfo {
    pub location_type: String,
    pub location: String,
    pub region: String,
    pub creds: AwsCredentials,
    // FIPS endpoint
    pub end_point: Option<String>,
}

#[derive(Deserialize, Debug)]
#[serde(rename_all = "SCREAMING_SNAKE_CASE")]
pub struct AwsCredentials {
    pub aws_key_id: String,
    pub aws_secret_key: String,
    pub aws_token: String,
    pub aws_id: String,
    pub aws_key: String,
}

#[derive(Deserialize, Debug)]
#[serde(rename_all = "camelCase")]
pub struct GcsPutGetStageInfo {
    pub location_type: String,
    pub location: String,
    pub storage_account: String,
    pub creds: GcsCredentials,
    pub presigned_url: String,
}

#[derive(Deserialize, Debug)]
#[serde(rename_all = "SCREAMING_SNAKE_CASE")]
pub struct GcsCredentials {
    pub gcs_access_token: String,
}

#[derive(Deserialize, Debug)]
#[serde(rename_all = "camelCase")]
pub struct AzurePutGetStageInfo {
    pub location_type: String,
    pub location: String,
    pub storage_account: String,
    pub creds: AzureCredentials,
}

#[derive(Deserialize, Debug)]
#[serde(rename_all = "SCREAMING_SNAKE_CASE")]
pub struct AzureCredentials {
    pub azure_sas_token: String,
}

#[derive(Deserialize, Debug)]
#[serde(untagged)]
pub enum EncryptionMaterialVariant {
    Single(PutGetEncryptionMaterial),
    Multiple(Vec<PutGetEncryptionMaterial>),
}

#[derive(Deserialize, Debug)]
#[serde(rename_all = "camelCase")]
pub struct PutGetEncryptionMaterial {
    // base64 encoded
    pub query_stage_master_key: String,
    pub query_id: String,
    pub smk_id: i64,
}