scout_rs 0.95.53

use std::collections::HashMap;

use anyhow::{anyhow, Result};
use postgrest::Postgrest;
use serde::{Deserialize, Serialize};

use crate::models::serde_helpers::serialized_object_key_set;
use crate::models::{Artifact, Event};

/// Response type for get_connectivity_for_artifact RPC (connectivity_with_coordinates).
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub struct ConnectivityWithCoordinates {
    pub id: Option<i64>,
    pub session_id: Option<i64>,
    pub device_id: Option<i64>,
    pub inserted_at: Option<String>,
    pub timestamp_start: Option<String>,
    pub signal: Option<f64>,
    pub noise: Option<f64>,
    pub altitude: Option<f64>,
    pub heading: Option<f64>,
    pub latitude: Option<f64>,
    pub longitude: Option<f64>,
    pub h14_index: Option<String>,
    pub h13_index: Option<String>,
    pub h12_index: Option<String>,
    pub h11_index: Option<String>,
    pub battery_percentage: Option<f32>,
    pub frequency_hz: Option<f32>,
    pub bandwidth_hz: Option<f32>,
    pub associated_station: Option<String>,
    pub mode: Option<String>,
    pub battery_volts: Option<f32>,
}


#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DatabaseConfig {
    pub rest_url: String,
    pub scout_api_key: String,
    pub supabase_api_key: String,
    /// When set, used as apikey + Bearer auth to bypass RLS; device api key is not sent.
    pub service_role_key: Option<String>,
    /// When true and service_role_key is set, connect() uses service key. When false, connect() uses device API key (required for identify/RPC).
    pub use_service_key: bool,
}

impl DatabaseConfig {
    /// Creates a new database config from environment variables with an optional Scout device API key
    /// If the scout device API key is not provided, it will fall back to the environment variable.
    ///
    /// The environment variables are:
    /// - SCOUT_DATABASE_REST_URL: The URL of the PostgREST API endpoint.
    /// - SCOUT_DEVICE_API_KEY: The API key for the Scout device (required when SCOUT_USE_SERVICE_KEY is not true).
    /// - SUPABASE_PUBLIC_API_KEY: The public API key for Supabase.
    /// - SUPABASE_SERVICE_ROLE_KEY: Optional. When set and SCOUT_USE_SERVICE_KEY=true, used to bypass RLS.
    /// - SCOUT_USE_SERVICE_KEY: Optional. When true and SUPABASE_SERVICE_ROLE_KEY is set, connect uses service key; otherwise connect uses device API key.
    pub fn from_env() -> Result<Self> {
        dotenv::dotenv().ok();

        let mut rest_url = std::env::var("SCOUT_DATABASE_REST_URL")
            .map_err(|_| anyhow!("SCOUT_DATABASE_REST_URL environment variable is required"))?;

        if !rest_url.ends_with("/rest/v1") {
            if rest_url.ends_with("/") {
                rest_url.push_str("rest/v1");
            } else {
                rest_url.push_str("/rest/v1");
            }
        }

        let service_role_key = std::env::var("SUPABASE_SERVICE_ROLE_KEY").ok();
        let use_service_key = std::env::var("SCOUT_USE_SERVICE_KEY")
            .map(|v| v == "1" || v.eq_ignore_ascii_case("true"))
            .unwrap_or(false);

        let (scout_api_key, supabase_api_key) = if use_service_key && service_role_key.is_some() {
            (std::env::var("SCOUT_DEVICE_API_KEY").unwrap_or_default(), std::env::var("SUPABASE_PUBLIC_API_KEY").unwrap_or_default())
        } else {
            let scout_api_key = std::env::var("SCOUT_DEVICE_API_KEY")
                .map_err(|_| anyhow!("SCOUT_DEVICE_API_KEY environment variable not found"))?;
            let supabase_api_key = std::env::var("SUPABASE_PUBLIC_API_KEY").map_err(|_| {
                anyhow!("SUPABASE_PUBLIC_API_KEY environment variable is required for Supabase access")
            })?;
            (scout_api_key, supabase_api_key)
        };

        Ok(DatabaseConfig {
            rest_url,
            scout_api_key,
            supabase_api_key,
            service_role_key,
            use_service_key,
        })
    }

    /// Service-role only config from env (for edge services: video-tagger, tracker).
    /// Needs SCOUT_DATABASE_REST_URL and SUPABASE_SERVICE_ROLE_KEY. Returns None if either is missing.
    pub fn from_env_service_role() -> Option<Self> {
        let rest_url = std::env::var("SCOUT_DATABASE_REST_URL").ok()?;
        let service_role_key = std::env::var("SUPABASE_SERVICE_ROLE_KEY").ok()?;
        let mut rest_url = rest_url;
        if !rest_url.ends_with("/rest/v1") {
            if rest_url.ends_with('/') {
                rest_url.push_str("rest/v1");
            } else {
                rest_url.push_str("/rest/v1");
            }
        }
        Some(DatabaseConfig {
            rest_url,
            scout_api_key: String::new(),
            supabase_api_key: String::new(),
            service_role_key: Some(service_role_key),
            use_service_key: true,
        })
    }

    /// Gets the PostgREST endpoint URL
    pub fn get_rest_url(&self) -> &str {
        &self.rest_url
    }

    /// Gets the Scout API key for custom authentication
    pub fn get_scout_api_key(&self) -> &str {
        &self.scout_api_key
    }

    /// Gets the Supabase API key for PostgREST access
    pub fn get_supabase_api_key(&self) -> &str {
        &self.supabase_api_key
    }
}

pub struct ScoutDbClient {
    config: DatabaseConfig,
    client: Option<Postgrest>,
}

impl std::fmt::Debug for ScoutDbClient {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("ScoutDbClient")
            .field("config", &self.config)
            .field(
                "client",
                if self.client.is_some() {
                    &"Connected"
                } else {
                    &"Disconnected"
                },
            )
            .finish()
    }
}

impl ScoutDbClient {
    pub fn new(config: DatabaseConfig) -> Self {
        Self {
            config,
            client: None,
        }
    }

    /// Establishes a connection to the database via PostgREST
    pub fn connect(&mut self) -> Result<()> {
        let rest_url = self.config.get_rest_url();

        let client = if self.config.use_service_key {
            if let Some(ref key) = self.config.service_role_key {
                Postgrest::new(rest_url)
                    .insert_header("apikey", key)
                    .insert_header("Authorization", format!("Bearer {}", key))
            } else {
                Postgrest::new(rest_url)
                    .insert_header("apikey", self.config.get_supabase_api_key())
                    .insert_header("api_key", &format!("{}", self.config.get_scout_api_key()))
            }
        } else {
            Postgrest::new(rest_url)
                .insert_header("apikey", self.config.get_supabase_api_key())
                .insert_header("api_key", &format!("{}", self.config.get_scout_api_key()))
        };

        self.client = Some(client);

        Ok(())
    }

    /// Gets the PostgREST client, ensuring connection is established
    pub fn get_client(&mut self) -> Result<&Postgrest> {
        if self.client.is_none() {
            self.connect()?;
        }

        self.client
            .as_ref()
            .ok_or_else(|| anyhow!("No PostgREST client available"))
    }

    /// Closes the database connection
    pub fn disconnect(&mut self) {
        if self.client.is_some() {
            self.client = None;
        }
    }

    /// Executes a query and returns the results
    pub async fn query<T>(
        &mut self,
        query_builder: impl FnOnce(&Postgrest) -> postgrest::Builder,
    ) -> Result<Vec<T>>
    where
        T: for<'de> serde::Deserialize<'de>,
    {
        let client = self.get_client()?;

        let builder = query_builder(client);
        let response = builder.execute().await?;

        let body = response.text().await?;

        // Try to parse as the expected type first
        if let Ok(results) = serde_json::from_str::<Vec<T>>(&body) {
            Ok(results)
        } else {
            // If that fails, try to parse as an error response
            if let Ok(error_response) = serde_json::from_str::<serde_json::Value>(&body) {
                if let Some(error_msg) = error_response.get("error") {
                    return Err(anyhow!("Database error: {}", error_msg));
                } else if let Some(message) = error_response.get("message") {
                    return Err(anyhow!("Database message: {}", message));
                } else {
                    return Err(anyhow!("Database returned unexpected format: {}", body));
                }
            } else {
                return Err(anyhow!(
                    "Failed to parse database response as JSON: {}",
                    body
                ));
            }
        }
    }

    /// Executes a query that returns a single row
    pub async fn query_one<T>(
        &mut self,
        query_builder: impl FnOnce(&Postgrest) -> postgrest::Builder,
    ) -> Result<T>
    where
        T: for<'de> serde::Deserialize<'de>,
    {
        let client = self.get_client()?;

        let builder = query_builder(client);
        let response = builder.execute().await?;

        let body = response.text().await?;
        let results: Vec<T> = serde_json::from_str(&body)?;

        if results.is_empty() {
            return Err(anyhow!("No results found"));
        }

        Ok(results.into_iter().next().unwrap())
    }

    /// Executes a query that doesn't return results (INSERT, UPDATE, DELETE)
    pub async fn execute(
        &mut self,
        query_builder: impl FnOnce(&Postgrest) -> postgrest::Builder,
    ) -> Result<()> {
        let client = self.get_client()?;

        let builder = query_builder(client);
        let response = builder.execute().await?;

        let status = response.status();
        if !status.is_success() {
            let error_text = response.text().await?;
            return Err(anyhow!(
                "Operation failed: HTTP {} - {}",
                status,
                error_text
            ));
        }

        Ok(())
    }

    /// Inserts data into a table
    pub async fn insert<T>(&mut self, table: &str, data: &T) -> Result<Vec<T>>
    where
        T: for<'de> serde::Deserialize<'de> + serde::Serialize,
    {
        let client = self.get_client()?;

        let json_data = serde_json::to_string(data)?;

        let response = client.from(table).insert(&json_data).execute().await?;

        let body = response.text().await?;

        // Try to parse as the expected type first
        if let Ok(results) = serde_json::from_str::<Vec<T>>(&body) {
            Ok(results)
        } else {
            // If that fails, try to parse as an error response
            if let Ok(error_response) = serde_json::from_str::<serde_json::Value>(&body) {
                if let Some(error_msg) = error_response.get("error") {
                    return Err(anyhow!("Database insert error: {}", error_msg));
                } else if let Some(message) = error_response.get("message") {
                    return Err(anyhow!("Database insert message: {}", message));
                } else {
                    return Err(anyhow!(
                        "Database insert returned unexpected format: {}",
                        body
                    ));
                }
            } else {
                return Err(anyhow!(
                    "Failed to parse database insert response as JSON: {}",
                    body
                ));
            }
        }
    }

    /// Inserts multiple items in a single bulk operation
    pub async fn insert_bulk<T>(&mut self, table: &str, data: &[T]) -> Result<Vec<T>>
    where
        T: for<'de> serde::Deserialize<'de> + serde::Serialize,
    {
        let client = self.get_client()?;

        let json_data = serde_json::to_string(data)?;

        let response = client.from(table).insert(&json_data).execute().await?;

        let body = response.text().await?;

        // Try to parse as the expected type first
        if let Ok(results) = serde_json::from_str::<Vec<T>>(&body) {
            Ok(results)
        } else {
            // If that fails, try to parse as an error response
            if let Ok(error_response) = serde_json::from_str::<serde_json::Value>(&body) {
                if let Some(error_msg) = error_response.get("error") {
                    return Err(anyhow!("Database bulk insert error: {}", error_msg));
                } else if let Some(message) = error_response.get("message") {
                    return Err(anyhow!("Database bulk insert message: {}", message));
                } else {
                    return Err(anyhow!(
                        "Database bulk insert returned unexpected format: {}",
                        body
                    ));
                }
            } else {
                return Err(anyhow!(
                    "Failed to parse database bulk insert response as JSON: {}",
                    body
                ));
            }
        }
    }

    /// Bulk insert split by serialized JSON key set (PostgREST needs uniform columns per request).
    pub async fn insert_bulk_batched_by_shape<T>(
        &mut self,
        table: &str,
        rows: &[T],
    ) -> Result<Vec<T>>
    where
        T: Serialize + Clone + for<'de> Deserialize<'de>,
    {
        if rows.is_empty() {
            return Ok(Vec::new());
        }
        let mut by_shape: HashMap<String, Vec<usize>> = HashMap::new();
        for (i, row) in rows.iter().enumerate() {
            let key = serialized_object_key_set(row)?;
            by_shape.entry(key).or_default().push(i);
        }
        let mut out = Vec::new();
        for indices in by_shape.values() {
            let batch: Vec<T> = indices.iter().map(|&i| rows[i].clone()).collect();
            let mut part = self.insert_bulk(table, &batch).await?;
            out.append(&mut part);
        }
        Ok(out)
    }

    /// Upserts multiple items in a single bulk operation (insert or update on conflict)
    pub async fn upsert_bulk<T>(&mut self, table: &str, data: &[T]) -> Result<Vec<T>>
    where
        T: for<'de> serde::Deserialize<'de> + serde::Serialize,
    {
        let client = self.get_client()?;

        let json_data = serde_json::to_string(data)?;

        let response = client
            .from(table)
            .upsert(&json_data)
            .on_conflict("id")
            .execute()
            .await?;

        let body = response.text().await?;

        // Try to parse as the expected type first
        if let Ok(results) = serde_json::from_str::<Vec<T>>(&body) {
            Ok(results)
        } else {
            // If that fails, try to parse as an error response
            if let Ok(error_response) = serde_json::from_str::<serde_json::Value>(&body) {
                if let Some(error_msg) = error_response.get("error") {
                    return Err(anyhow!("Database bulk upsert error: {}", error_msg));
                } else if let Some(message) = error_response.get("message") {
                    return Err(anyhow!("Database bulk upsert message: {}", message));
                } else {
                    return Err(anyhow!(
                        "Database bulk upsert returned unexpected format: {}",
                        body
                    ));
                }
            } else {
                return Err(anyhow!(
                    "Failed to parse database bulk upsert response as JSON: {}",
                    body
                ));
            }
        }
    }

    /// Bulk upsert split by serialized JSON key set (PostgREST needs uniform columns per request).
    pub async fn upsert_bulk_batched_by_shape<T>(
        &mut self,
        table: &str,
        rows: &[T],
    ) -> Result<Vec<T>>
    where
        T: Serialize + Clone + for<'de> Deserialize<'de>,
    {
        if rows.is_empty() {
            return Ok(Vec::new());
        }
        let mut by_shape: HashMap<String, Vec<usize>> = HashMap::new();
        for (i, row) in rows.iter().enumerate() {
            let key = serialized_object_key_set(row)?;
            by_shape.entry(key).or_default().push(i);
        }
        let mut out = Vec::new();
        for indices in by_shape.values() {
            let batch: Vec<T> = indices.iter().map(|&i| rows[i].clone()).collect();
            let mut part = self.upsert_bulk(table, &batch).await?;
            out.append(&mut part);
        }
        Ok(out)
    }

    /// Updates data in a table
    pub async fn update<T>(
        &mut self,
        data: &T,
        filter_builder: impl FnOnce(&Postgrest) -> postgrest::Builder,
    ) -> Result<Vec<T>>
    where
        T: for<'de> serde::Deserialize<'de> + serde::Serialize,
    {
        let client = self.get_client()?;

        let json_data = serde_json::to_string(data)?;

        let builder = filter_builder(client);
        let response = builder.update(&json_data).execute().await?;

        let status = response.status();
        let body = response.text().await?;

        if !status.is_success() {
            if let Ok(v) = serde_json::from_str::<serde_json::Value>(&body) {
                if let Some(m) = v.get("message").and_then(|x| x.as_str()) {
                    return Err(anyhow!("Database update error: {} (HTTP {})", m, status));
                }
            }
            return Err(anyhow!(
                "Database update failed: HTTP {} - {}",
                status,
                body
            ));
        }

        if body.trim().is_empty() {
            return Ok(vec![]);
        }

        if let Ok(results) = serde_json::from_str::<Vec<T>>(&body) {
            return Ok(results);
        }

        if let Ok(error_response) = serde_json::from_str::<serde_json::Value>(&body) {
            if let Some(error_msg) = error_response.get("error") {
                return Err(anyhow!("Database update error: {}", error_msg));
            } else if let Some(message) = error_response.get("message") {
                return Err(anyhow!("Database update message: {}", message));
            }
            return Err(anyhow!(
                "Database update returned unexpected format: {}",
                body
            ));
        }

        Err(anyhow!(
            "Failed to parse database update response as JSON: {}",
            body
        ))
    }

    /// Deletes data from a table
    pub async fn delete(
        &mut self,
        filter_builder: impl FnOnce(&Postgrest) -> postgrest::Builder,
    ) -> Result<()> {
        let client = self.get_client()?;

        let builder = filter_builder(client);
        let response = builder.delete().execute().await?;

        let status = response.status();
        if !status.is_success() {
            let error_text = response.text().await?;
            return Err(anyhow!(
                "Delete operation failed: HTTP {} - {}",
                status,
                error_text
            ));
        }

        Ok(())
    }

    /// Fetches events for the given primary key ids. Returns an empty vector when `ids` is empty.
    pub async fn get_events_by_ids(&mut self, ids: &[i64]) -> Result<Vec<Event>> {
        if ids.is_empty() {
            return Ok(vec![]);
        }
        let id_strings: Vec<String> = ids.iter().map(|id| id.to_string()).collect();
        self.query(move |client| client.from("events").select("*").in_("id", id_strings))
            .await
    }

    /// Fetches artifacts for the given primary key ids. Returns an empty vector when `ids` is empty.
    pub async fn get_artifacts_by_ids(&mut self, ids: &[i64]) -> Result<Vec<Artifact>> {
        if ids.is_empty() {
            return Ok(vec![]);
        }
        let id_strings: Vec<String> = ids.iter().map(|id| id.to_string()).collect();
        self.query(move |client| client.from("artifacts").select("*").in_("id", id_strings))
            .await
    }

    /// Connectivity records for an artifact (session or device+time window). Uses RPC get_connectivity_for_artifact.
    pub async fn get_connectivity_for_artifact(
        &mut self,
        artifact_id: i64,
        max_elements: Option<i32>,
    ) -> Result<Vec<ConnectivityWithCoordinates>> {
        let client = self.get_client()?;
        let max = max_elements.unwrap_or(1000);
        let response = client
            .rpc(
                "get_connectivity_for_artifact",
                serde_json::json!({
                    "artifact_id_caller": artifact_id,
                    "max_elements_caller": max
                })
                .to_string(),
            )
            .execute()
            .await?;

        let status = response.status();
        let body = response.text().await?;
        if !status.is_success() {
            return Err(anyhow!("get_connectivity_for_artifact failed: HTTP {} - {}", status, body));
        }

        let results: Vec<ConnectivityWithCoordinates> = serde_json::from_str(&body)
            .map_err(|e| anyhow!("Failed to parse get_connectivity_for_artifact response: {} - {}", e, body))?;
        Ok(results)
    }

    /// Connectivity records for an event (device + time window around timestamp_observation). Uses RPC get_connectivity_for_event.
    pub async fn get_connectivity_for_event(
        &mut self,
        event_id: i64,
        max_elements: Option<i32>,
    ) -> Result<Vec<ConnectivityWithCoordinates>> {
        let client = self.get_client()?;
        let max = max_elements.unwrap_or(1000);
        let response = client
            .rpc(
                "get_connectivity_for_event",
                serde_json::json!({
                    "event_id_caller": event_id,
                    "max_elements_caller": max
                })
                .to_string(),
            )
            .execute()
            .await?;

        let status = response.status();
        let body = response.text().await?;
        if !status.is_success() {
            return Err(anyhow!("get_connectivity_for_event failed: HTTP {} - {}", status, body));
        }

        let results: Vec<ConnectivityWithCoordinates> = serde_json::from_str(&body)
            .map_err(|e| anyhow!("Failed to parse get_connectivity_for_event response: {} - {}", e, body))?;
        Ok(results)
    }

    /// Deletes a message from the named pgmq queue (calls public.ack_queue_message).
    /// Used by the GCP video-tagger after processing an embed or tag job.
    pub async fn ack_queue_message(&mut self, queue_name: &str, message_id: i64) -> Result<()> {
        let client = self.get_client()?;
        let response = client
            .rpc(
                "ack_queue_message",
                serde_json::json!({
                    "queue_name": queue_name,
                    "message_id": message_id
                })
                .to_string(),
            )
            .execute()
            .await?;

        let status = response.status();
        if !status.is_success() {
            let body = response.text().await?;
            return Err(anyhow!(
                "ack_queue_message failed: HTTP {} - {}",
                status,
                body
            ));
        }
        Ok(())
    }
}

impl Drop for ScoutDbClient {
    fn drop(&mut self) {
        self.disconnect();
    }
}