use anyhow::Context;
use axum::{
    extract::{Path, Query},
    routing::{delete, get, post},
    Extension, Json, Router,
};
use futures::{stream::SplitSink, SinkExt, StreamExt};
use http::StatusCode;
use itertools::Itertools;
use rand::seq::SliceRandom;
use serde::{
    de::{self, MapAccess, Visitor},
    Deserialize, Deserializer, Serialize,
};
use serde_json::{value::RawValue, Value};
use sql_builder::{bind::Bind, SqlBuilder};
use sqlx::prelude::FromRow;
use sqlx::types::Json as SqlxJson;
use std::{collections::HashMap, fmt};
use tokio::net::TcpStream;
use tokio_tungstenite::{connect_async, tungstenite::Message, MaybeTlsStream, WebSocketStream};
use uuid::Uuid;
use windmill_audit::{audit_ee::audit_log, ActionKind};
use windmill_common::{
    db::UserDB,
    error::{self, to_anyhow, JsonResult},
    utils::{not_found_if_none, paginate, report_critical_error, Pagination, StripPath},
    worker::{to_raw_value, CLOUD_HOSTED},
    INSTANCE_NAME,
};
use windmill_queue::PushArgsOwned;

use crate::{
    capture::{insert_capture_payload, TriggerKind, WebsocketTriggerConfig},
    db::{ApiAuthed, DB},
    jobs::{run_flow_by_path_inner, run_script_by_path_inner, RunJobQuery},
    users::fetch_api_authed,
};

use std::borrow::Cow;

pub fn workspaced_service() -> Router {
    Router::new()
        .route("/create", post(create_websocket_trigger))
        .route("/list", get(list_websocket_triggers))
        .route("/get/*path", get(get_websocket_trigger))
        .route("/update/*path", post(update_websocket_trigger))
        .route("/delete/*path", delete(delete_websocket_trigger))
        .route("/exists/*path", get(exists_websocket_trigger))
        .route("/setenabled/*path", post(set_enabled))
}

#[derive(Deserialize)]
struct NewWebsocketTrigger {
    path: String,
    url: String,
    script_path: String,
    is_flow: bool,
    enabled: Option<bool>,
    filters: Vec<Box<RawValue>>,
    initial_messages: Option<Vec<Box<RawValue>>>,
    url_runnable_args: Option<Box<RawValue>>,
}

#[derive(Deserialize)]
pub struct JsonFilter {
    key: String,
    value: serde_json::Value,
}

#[derive(Deserialize)]
#[serde(untagged)]
pub enum Filter {
    JsonFilter(JsonFilter),
}

#[derive(Deserialize)]
enum InitialMessage {
    #[serde(rename = "raw_message")]
    RawMessage(String),
    #[serde(rename = "runnable_result")]
    RunnableResult { path: String, args: Box<RawValue>, is_flow: bool },
}

#[derive(FromRow, Serialize, Clone)]
pub struct WebsocketTrigger {
    workspace_id: String,
    path: String,
    url: String,
    script_path: String,
    is_flow: bool,
    edited_by: String,
    email: String,
    edited_at: chrono::DateTime<chrono::Utc>,
    server_id: Option<String>,
    last_server_ping: Option<chrono::DateTime<chrono::Utc>>,
    extra_perms: serde_json::Value,
    error: Option<String>,
    enabled: bool,
    filters: Vec<SqlxJson<Box<RawValue>>>,
    initial_messages: Option<Vec<SqlxJson<Box<RawValue>>>>,
    url_runnable_args: Option<SqlxJson<Box<RawValue>>>,
}

#[derive(Deserialize)]
struct EditWebsocketTrigger {
    path: String,
    url: String,
    script_path: String,
    is_flow: bool,
    filters: Vec<Box<RawValue>>,
    initial_messages: Option<Vec<Box<RawValue>>>,
    url_runnable_args: Option<Box<RawValue>>,
}

#[derive(Deserialize)]
pub struct ListWebsocketTriggerQuery {
    pub page: Option<usize>,
    pub per_page: Option<usize>,
    pub path: Option<String>,
    pub is_flow: Option<bool>,
    pub path_start: Option<String>,
}

async fn list_websocket_triggers(
    authed: ApiAuthed,
    Extension(user_db): Extension<UserDB>,
    Path(w_id): Path<String>,
    Query(lst): Query<ListWebsocketTriggerQuery>,
) -> error::JsonResult<Vec<WebsocketTrigger>> {
    let mut tx = user_db.begin(&authed).await?;
    let (per_page, offset) = paginate(Pagination { per_page: lst.per_page, page: lst.page });
    let mut sqlb = SqlBuilder::select_from("websocket_trigger")
        .field("*")
        .order_by("edited_at", true)
        .and_where("workspace_id = ?".bind(&w_id))
        .offset(offset)
        .limit(per_page)
        .clone();
    if let Some(path) = lst.path {
        sqlb.and_where_eq("script_path", "?".bind(&path));
    }
    if let Some(is_flow) = lst.is_flow {
        sqlb.and_where_eq("is_flow", "?".bind(&is_flow));
    }
    if let Some(path_start) = &lst.path_start {
        sqlb.and_where_like_left("path", path_start);
    }
    let sql = sqlb
        .sql()
        .map_err(|e| error::Error::InternalErr(e.to_string()))?;
    let rows = sqlx::query_as::<_, WebsocketTrigger>(&sql)
        .fetch_all(&mut *tx)
        .await?;
    tx.commit().await?;

    Ok(Json(rows))
}

async fn get_websocket_trigger(
    authed: ApiAuthed,
    Extension(user_db): Extension<UserDB>,
    Path((w_id, path)): Path<(String, StripPath)>,
) -> error::JsonResult<WebsocketTrigger> {
    let mut tx = user_db.begin(&authed).await?;
    let path = path.to_path();
    let trigger = sqlx::query_as::<_, WebsocketTrigger>(
        r#"SELECT *
          FROM websocket_trigger
          WHERE workspace_id = $1 AND path = $2"#,
    )
    .bind(w_id)
    .bind(path)
    .fetch_optional(&mut *tx)
    .await?;
    tx.commit().await?;

    let trigger = not_found_if_none(trigger, "Trigger", path)?;

    Ok(Json(trigger))
}

async fn create_websocket_trigger(
    authed: ApiAuthed,
    Extension(user_db): Extension<UserDB>,
    Path(w_id): Path<String>,
    Json(ct): Json<NewWebsocketTrigger>,
) -> error::Result<(StatusCode, String)> {
    if *CLOUD_HOSTED {
        return Err(error::Error::BadRequest(
            "Websocket triggers are not supported on multi-tenant cloud, use dedicated cloud or self-host".to_string(),
        ));
    }

    let mut tx = user_db.begin(&authed).await?;

    let filters = ct.filters.into_iter().map(SqlxJson).collect_vec();
    let initial_messages = ct
        .initial_messages
        .unwrap_or_default()
        .into_iter()
        .map(SqlxJson)
        .collect_vec();
    sqlx::query_as::<_, WebsocketTrigger>(
      "INSERT INTO websocket_trigger (workspace_id, path, url, script_path, is_flow, enabled, filters, initial_messages, url_runnable_args, edited_by, email, edited_at) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, now()) RETURNING *",
    )
    .bind(&w_id)
    .bind(&ct.path)
    .bind(ct.url)
    .bind(ct.script_path)
    .bind(ct.is_flow)
    .bind(ct.enabled.unwrap_or(true))
    .bind(filters.as_slice())
    .bind(initial_messages.as_slice())
    .bind(ct.url_runnable_args.map(SqlxJson))
    .bind(&authed.username)
    .bind(&authed.email)
    .fetch_one(&mut *tx).await?;

    audit_log(
        &mut *tx,
        &authed,
        "websocket_triggers.create",
        ActionKind::Create,
        &w_id,
        Some(ct.path.as_str()),
        None,
    )
    .await?;

    tx.commit().await?;

    Ok((StatusCode::CREATED, format!("{}", ct.path)))
}

async fn update_websocket_trigger(
    authed: ApiAuthed,
    Extension(user_db): Extension<UserDB>,
    Path((w_id, path)): Path<(String, StripPath)>,
    Json(ct): Json<EditWebsocketTrigger>,
) -> error::Result<String> {
    let path = path.to_path();
    let mut tx = user_db.begin(&authed).await?;

    let filters = ct.filters.into_iter().map(SqlxJson).collect_vec();
    let initial_messages = ct
        .initial_messages
        .unwrap_or_default()
        .into_iter()
        .map(SqlxJson)
        .collect_vec();

    // important to update server_id, last_server_ping and error to NULL to stop current websocket listener
    sqlx::query!(
        "UPDATE websocket_trigger SET url = $1, script_path = $2, path = $3, is_flow = $4, filters = $5, initial_messages = $6, url_runnable_args = $7, edited_by = $8, email = $9, edited_at = now(), server_id = NULL, error = NULL
            WHERE workspace_id = $10 AND path = $11",
        ct.url,
        ct.script_path,
        ct.path,
        ct.is_flow,
        filters.as_slice() as &[SqlxJson<Box<RawValue>>],
        initial_messages.as_slice() as &[SqlxJson<Box<RawValue>>],
        ct.url_runnable_args.map(SqlxJson) as Option<SqlxJson<Box<RawValue>>>,
        &authed.username,
        &authed.email,
        w_id,
        path,
    )
    .execute(&mut *tx).await?;

    audit_log(
        &mut *tx,
        &authed,
        "websocket_triggers.update",
        ActionKind::Create,
        &w_id,
        Some(path),
        None,
    )
    .await?;

    tx.commit().await?;

    Ok(path.to_string())
}

#[derive(Deserialize)]
pub struct SetEnabled {
    pub enabled: bool,
}

pub async fn set_enabled(
    authed: ApiAuthed,
    Extension(user_db): Extension<UserDB>,
    Path((w_id, path)): Path<(String, StripPath)>,
    Json(payload): Json<SetEnabled>,
) -> error::Result<String> {
    let mut tx = user_db.begin(&authed).await?;
    let path = path.to_path();

    // important to set server_id, last_server_ping and error to NULL to stop current websocket listener
    let one_o = sqlx::query_scalar!(
        "UPDATE websocket_trigger SET enabled = $1, email = $2, edited_by = $3, edited_at = now(), server_id = NULL, error = NULL
        WHERE path = $4 AND workspace_id = $5 RETURNING 1",
        payload.enabled,
        &authed.email,
        &authed.username,
        path,
        w_id,
    ).fetch_optional(&mut *tx).await?;

    not_found_if_none(one_o.flatten(), "Websocket trigger", path)?;

    audit_log(
        &mut *tx,
        &authed,
        "websocket_triggers.setenabled",
        ActionKind::Update,
        &w_id,
        Some(path),
        Some([("enabled", payload.enabled.to_string().as_ref())].into()),
    )
    .await?;

    tx.commit().await?;

    Ok(format!(
        "succesfully updated websocket trigger at path {} to status {}",
        path, payload.enabled
    ))
}

async fn delete_websocket_trigger(
    authed: ApiAuthed,
    Extension(user_db): Extension<UserDB>,
    Path((w_id, path)): Path<(String, StripPath)>,
) -> error::Result<String> {
    let path = path.to_path();
    let mut tx = user_db.begin(&authed).await?;
    sqlx::query!(
        "DELETE FROM websocket_trigger WHERE workspace_id = $1 AND path = $2",
        w_id,
        path,
    )
    .execute(&mut *tx)
    .await?;

    audit_log(
        &mut *tx,
        &authed,
        "websocket_triggers.delete",
        ActionKind::Delete,
        &w_id,
        Some(path),
        None,
    )
    .await?;

    tx.commit().await?;

    Ok(format!("Websocket trigger {path} deleted"))
}

async fn exists_websocket_trigger(
    Extension(db): Extension<DB>,
    Path((w_id, path)): Path<(String, StripPath)>,
) -> JsonResult<bool> {
    let path = path.to_path();
    let exists = sqlx::query_scalar!(
        "SELECT EXISTS(SELECT 1 FROM websocket_trigger WHERE path = $1 AND workspace_id = $2)",
        path,
        w_id,
    )
    .fetch_one(&db)
    .await?
    .unwrap_or(false);
    Ok(Json(exists))
}

async fn listen_to_unlistened_websockets(
    db: &DB,
    killpill_rx: &tokio::sync::broadcast::Receiver<()>,
) -> () {
    match sqlx::query_as::<_, WebsocketTrigger>(
        r#"SELECT *
            FROM websocket_trigger
            WHERE enabled IS TRUE AND (last_server_ping IS NULL OR last_server_ping < now() - interval '15 seconds')"#
    )
    .fetch_all(db)
    .await
    {
        Ok(mut triggers) => {
            triggers.shuffle(&mut rand::thread_rng());
            for trigger in triggers {
                trigger.maybe_listen_to_websocket(db.clone(), killpill_rx.resubscribe()).await;
            }
        }
        Err(err) => {
            tracing::error!("Error fetching websocket triggers: {:?}", err);
        }
    };

    match sqlx::query_as!(
        CaptureConfigForWebsocket,
        r#"SELECT path, is_flow, workspace_id, trigger_config as "trigger_config!: _", owner, email FROM capture_config WHERE trigger_kind = 'websocket' AND last_client_ping > NOW() - INTERVAL '10 seconds' AND trigger_config IS NOT NULL AND (last_server_ping IS NULL OR last_server_ping < now() - interval '15 seconds')"#
    )
    .fetch_all(db)
    .await
    {
        Ok(mut captures) => {
            captures.shuffle(&mut rand::thread_rng());
            for capture in captures {
                capture.maybe_listen_to_websocket(db.clone(),  killpill_rx.resubscribe()).await;
            }
        }
        Err(err) => {
            tracing::error!("Error fetching capture websocket triggers: {:?}", err);
        }
    }
}

pub async fn start_websockets(db: DB, mut killpill_rx: tokio::sync::broadcast::Receiver<()>) -> () {
    tokio::spawn(async move {
        listen_to_unlistened_websockets(&db, &killpill_rx).await;
        loop {
            tokio::select! {
                biased;
                _ = killpill_rx.recv() => {
                    return;
                }
                _ = tokio::time::sleep(tokio::time::Duration::from_secs(15)) => {
                    listen_to_unlistened_websockets(&db, &killpill_rx).await;
                }
            }
        }
    });
}

struct SupersetVisitor<'a> {
    key: &'a str,
    value_to_check: &'a Value,
}

impl<'de, 'a> Visitor<'de> for SupersetVisitor<'a> {
    type Value = bool;

    fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
        formatter.write_str("a JSON object with a specific key at the top level")
    }

    fn visit_map<V>(self, mut map: V) -> Result<Self::Value, V::Error>
    where
        V: MapAccess<'de>,
    {
        while let Some(key) = map.next_key::<String>()? {
            if key == self.key {
                // Deserialize the value for the key and check if it's a superset
                let json_value: Value = map.next_value()?;
                return Ok(is_superset(&json_value, self.value_to_check));
            } else {
                // Skip the value if it's not the one we're interested in
                let _ = map.next_value::<de::IgnoredAny>()?;
            }
        }
        // If the key was not found, return false
        Ok(false)
    }
}

// Function to check if json_value is a superset of value_to_check
fn is_superset(json_value: &Value, value_to_check: &Value) -> bool {
    match (json_value, value_to_check) {
        (Value::Object(json_map), Value::Object(check_map)) => {
            // Check that all keys and values in check_map exist and match in json_map
            check_map.iter().all(|(k, v)| {
                json_map
                    .get(k)
                    .map_or(false, |json_val| is_superset(json_val, v))
            })
        }
        (Value::Array(json_array), Value::Array(check_array)) => {
            // Check that all elements in check_array exist in json_array
            check_array.iter().all(|check_item| {
                json_array
                    .iter()
                    .any(|json_item| is_superset(json_item, check_item))
            })
        }
        _ => json_value == value_to_check,
    }
}

// A function to deserialize and check if the value at the given key is a superset of a passed value
fn is_value_superset<'a, 'de, D>(
    deserializer: D,
    key: &'a str,
    value_to_check: &'a Value,
) -> Result<bool, D::Error>
where
    D: Deserializer<'de>,
{
    deserializer.deserialize_map(SupersetVisitor { key, value_to_check })
}

async fn wait_runnable_result(
    path: String,
    is_flow: bool,
    args: Option<&Box<RawValue>>,
    authed: ApiAuthed,
    db: &DB,
    workspace_id: &str,
    trigger_path: &str,
) -> error::Result<String> {
    let user_db = UserDB::new(db.clone());

    let args = if let Some(args) = args {
        serde_json::from_str::<Option<HashMap<String, Box<RawValue>>>>(args.get())
            .map_err(|e| error::Error::BadRequest(format!("invalid json: {}", e)))?
            .unwrap_or_else(HashMap::new)
    } else {
        HashMap::new()
    };

    let label_prefix = Some(format!("ws-{}-", trigger_path));
    let (_, job_id) = if is_flow {
        run_flow_by_path_inner(
            authed,
            db.clone(),
            user_db,
            workspace_id.to_string(),
            StripPath(path.clone()),
            RunJobQuery::default(),
            PushArgsOwned { args, extra: None },
            label_prefix,
        )
        .await?
    } else {
        run_script_by_path_inner(
            authed,
            db.clone(),
            user_db,
            workspace_id.to_string(),
            StripPath(path.clone()),
            RunJobQuery::default(),
            PushArgsOwned { args, extra: None },
            label_prefix,
        )
        .await?
    };

    let start_time = tokio::time::Instant::now();

    loop {
        if start_time.elapsed() > tokio::time::Duration::from_secs(300) {
            return Err(anyhow::anyhow!(
                "Timed out after 5m waiting for runnable {path} (is_flow: {is_flow}) to complete",
            )
            .into());
        }

        #[derive(sqlx::FromRow)]
        struct RawResult {
            result: Option<SqlxJson<Box<RawValue>>>,
            success: bool,
        }

        let result = sqlx::query_as::<_, RawResult>(
            "SELECT result, success FROM completed_job WHERE id = $1 AND workspace_id = $2",
        )
        .bind(Uuid::parse_str(&job_id).unwrap())
        .bind(workspace_id)
        .fetch_optional(db)
        .await;

        match result {
            Ok(Some(r)) => {
                if !r.success {
                    return Err(anyhow::anyhow!(
                        "Runnable {path} (is_flow: {is_flow}) failed: {:?}",
                        r.result
                    )
                    .into());
                } else {
                    return Ok(r.result.map(|r| r.get().to_owned()).unwrap_or_default());
                }
            }
            Ok(None) => {
                // not yet done, wait for 5s and check again
                tokio::time::sleep(tokio::time::Duration::from_secs(5)).await;
            }
            Err(err) => {
                return Err(anyhow::anyhow!(
                    "Error fetching job result for runnable {path} (is_flow: {is_flow}): {err}",
                )
                .into());
            }
        }
    }
}

async fn loop_ping(db: &DB, ws: &WebsocketEnum, error: Option<&str>) -> () {
    loop {
        if let None = ws.update_ping(db, error).await {
            return;
        }
        tokio::time::sleep(tokio::time::Duration::from_secs(5)).await;
    }
}

async fn get_url_from_runnable(
    path: &str,
    is_flow: bool,
    db: &DB,
    authed: ApiAuthed,
    args: Option<&Box<RawValue>>,
    workspace_id: &str,
    trigger_path: &str,
) -> error::Result<String> {
    tracing::info!("Running runnable {path} (is_flow: {is_flow}) to get websocket URL",);

    let result = wait_runnable_result(
        path.to_string(),
        is_flow,
        args,
        authed,
        db,
        workspace_id,
        trigger_path,
    )
    .await?;

    if result.starts_with("\"") && result.ends_with("\"") {
        Ok(result[1..result.len() - 1].to_string())
    } else {
        Err(anyhow::anyhow!("Runnable {path} (is_flow: {is_flow}) did not return a string").into())
    }
}

impl WebsocketTrigger {
    async fn maybe_listen_to_websocket(
        self,
        db: DB,
        killpill_rx: tokio::sync::broadcast::Receiver<()>,
    ) -> () {
        match sqlx::query_scalar!(
            "UPDATE websocket_trigger SET server_id = $1, last_server_ping = now(), error = 'Connecting...' WHERE enabled IS TRUE AND workspace_id = $2 AND path = $3 AND (last_server_ping IS NULL OR last_server_ping < now() - interval '15 seconds') RETURNING true",
            *INSTANCE_NAME,
            self.workspace_id,
            self.path,
        ).fetch_optional(&db).await {
            Ok(has_lock) => {
                if has_lock.flatten().unwrap_or(false) {
                    tokio::spawn(listen_to_websocket(WebsocketEnum::Trigger(self), db, killpill_rx));
                } else {
                    tracing::info!("Websocket {} already being listened to", self.url);
                }
            },
            Err(err) => {
                tracing::error!("Error acquiring lock for websocket {}: {:?}", self.path, err);
            }
        };
    }

    async fn update_ping(&self, db: &DB, error: Option<&str>) -> Option<()> {
        match sqlx::query_scalar!(
        "UPDATE websocket_trigger SET last_server_ping = now(), error = $1 WHERE workspace_id = $2 AND path = $3 AND server_id = $4 AND enabled IS TRUE RETURNING 1",
        error,
        self.workspace_id,
        self.path,
        *INSTANCE_NAME
    ).fetch_optional(db).await {
        Ok(updated) => {
            if updated.flatten().is_none() {
                tracing::info!("Websocket {} changed, disabled, or deleted, stopping...", self.url); 
                return None;
            }
        },
            Err(err) => {
                tracing::warn!("Error updating ping of websocket {}: {:?}", self.url, err);
            }
        };

        Some(())
    }

    async fn disable_with_error(&self, db: &DB, error: String) -> () {
        match sqlx::query!(
            "UPDATE websocket_trigger SET enabled = FALSE, error = $1, server_id = NULL, last_server_ping = NULL WHERE workspace_id = $2 AND path = $3",
            error,
            self.workspace_id,
            self.path,
        )
        .execute(db).await {
            Ok(_) => {
                report_critical_error(format!("Disabling websocket {} because of error: {}", self.url, error), db.clone(), Some(&self.workspace_id), None).await;
            },
            Err(disable_err) => {
                report_critical_error(
                    format!("Could not disable websocket {} with err {}, disabling because of error {}", self.path, disable_err, error), 
                    db.clone(),
                    Some(&self.workspace_id),
                    None,
                ).await;
            }
        }
    }

    async fn get_url_from_runnable(
        &self,
        path: &str,
        is_flow: bool,
        db: &DB,
    ) -> error::Result<String> {
        get_url_from_runnable(
            &path,
            is_flow,
            db,
            self.fetch_authed(db, Some("url".to_string())).await?,
            self.url_runnable_args.as_ref().map(|r| &r.0),
            &self.workspace_id,
            &self.path,
        )
        .await
    }

    async fn send_initial_messages(
        &self,
        mut writer: SplitSink<WebSocketStream<MaybeTlsStream<TcpStream>>, Message>,
        db: &DB,
    ) -> error::Result<()> {
        let initial_messages: Vec<InitialMessage> = self
            .initial_messages
            .as_deref()
            .unwrap_or_default()
            .iter()
            .filter_map(|m| serde_json::from_str(m.get()).ok())
            .collect_vec();

        for start_message in initial_messages {
            match start_message {
                InitialMessage::RawMessage(msg) => {
                    let msg = if msg.starts_with("\"") && msg.ends_with("\"") {
                        msg[1..msg.len() - 1].to_string()
                    } else {
                        msg
                    };
                    tracing::info!(
                        "Sending raw message initial message to websocket {}: {}",
                        self.url,
                        msg
                    );
                    writer
                        .send(tokio_tungstenite::tungstenite::Message::Text(msg))
                        .await
                        .map_err(to_anyhow)
                        .with_context(|| "failed to send raw message")?;
                }
                InitialMessage::RunnableResult { path, is_flow, args } => {
                    tracing::info!(
                        "Running runnable {path} (is_flow: {is_flow}) for initial message to websocket {}",
                        self.url,
                    );

                    let result = wait_runnable_result(
                        path.clone(),
                        is_flow,
                        Some(&args),
                        self.fetch_authed(db, Some("init".to_string())).await?,
                        db,
                        &self.workspace_id,
                        &self.path,
                    )
                    .await?;

                    tracing::info!(
                        "Sending runnable {path} (is_flow: {is_flow}) result to websocket {}",
                        self.url
                    );

                    let result = if result.starts_with("\"") && result.ends_with("\"") {
                        result[1..result.len() - 1].to_string()
                    } else {
                        result
                    };

                    writer
                        .send(tokio_tungstenite::tungstenite::Message::Text(result))
                        .await
                        .map_err(to_anyhow)
                        .with_context(|| {
                            format!("Failed to send runnable {path} (is_flow: {is_flow}) result")
                        })?;
                }
            }
        }

        Ok(())
    }

    async fn handle(&self, db: &DB, args: PushArgsOwned) -> () {
        if let Err(err) = run_job(db, self, args).await {
            report_critical_error(
                format!(
                    "Failed to trigger job from websocket {}: {:?}",
                    self.url, err
                ),
                db.clone(),
                Some(&self.workspace_id),
                None,
            )
            .await;
        };
    }

    async fn fetch_authed(
        &self,
        db: &DB,
        username_override: Option<String>,
    ) -> error::Result<ApiAuthed> {
        fetch_api_authed(
            self.edited_by.clone(),
            self.email.clone(),
            &self.workspace_id,
            db,
            username_override,
        )
        .await
    }
}

#[derive(Deserialize)]
struct CaptureConfigForWebsocket {
    trigger_config: SqlxJson<WebsocketTriggerConfig>,
    path: String,
    is_flow: bool,
    workspace_id: String,
    owner: String,
    email: String,
}

impl CaptureConfigForWebsocket {
    async fn maybe_listen_to_websocket(
        self,
        db: DB,
        killpill_rx: tokio::sync::broadcast::Receiver<()>,
    ) -> () {
        match sqlx::query_scalar!(
            "UPDATE capture_config SET server_id = $1, last_server_ping = now(), error = 'Connecting...' WHERE last_client_ping > NOW() - INTERVAL '10 seconds' AND workspace_id = $2 AND path = $3 AND is_flow = $4 AND trigger_kind = 'websocket' AND (last_server_ping IS NULL OR last_server_ping < now() - interval '15 seconds') RETURNING true",
            *INSTANCE_NAME,
            self.workspace_id,
            self.path,
            self.is_flow,
        ).fetch_optional(&db).await {
            Ok(has_lock) => {
                if has_lock.flatten().unwrap_or(false) {
                    tokio::spawn(listen_to_websocket(WebsocketEnum::Capture(self), db, killpill_rx));
                } else {
                    tracing::info!("Websocket {} already being listened to", self.trigger_config.url);
                }
            },
            Err(err) => {
                tracing::error!("Error acquiring lock for capture websocket {}: {:?}", self.path, err);
            }
        };
    }

    async fn update_ping(&self, db: &DB, error: Option<&str>) -> Option<()> {
        match sqlx::query_scalar!(
        "UPDATE capture_config SET last_server_ping = now(), error = $1 WHERE workspace_id = $2 AND path = $3 AND is_flow = $4 AND trigger_kind = 'websocket' AND server_id = $5 AND last_client_ping > NOW() - INTERVAL '10 seconds' RETURNING 1",
        error,
        self.workspace_id,
        self.path,
        self.is_flow,
        *INSTANCE_NAME
    ).fetch_optional(db).await {
        Ok(updated) => {
            if updated.flatten().is_none() {
                tracing::info!("Websocket capture {} changed, disabled, or deleted, stopping...", self.trigger_config.url); 
                return None;
            }
        },
            Err(err) => {
                tracing::warn!("Error updating ping of capture websocket {}: {:?}", self.trigger_config.url, err);
            }
        };

        Some(())
    }

    async fn handle(&self, db: &DB, args: PushArgsOwned) -> () {
        if let Err(err) = insert_capture_payload(
            db,
            &self.workspace_id,
            &self.path,
            self.is_flow,
            &TriggerKind::Websocket,
            PushArgsOwned { args: args.args, extra: None },
            args.extra.as_ref().map(to_raw_value),
            &self.owner,
        )
        .await
        {
            tracing::error!("Error inserting capture payload: {:?}", err);
        }
    }

    async fn get_url_from_runnable(
        &self,
        path: &str,
        is_flow: bool,
        db: &DB,
    ) -> error::Result<String> {
        let url_runnable_args = self
            .trigger_config
            .url_runnable_args
            .as_ref()
            .map(to_raw_value);
        get_url_from_runnable(
            &path,
            is_flow,
            db,
            self.fetch_authed(db, Some("url".to_string())).await?,
            url_runnable_args.as_ref(),
            &self.workspace_id,
            &self.get_trigger_path(),
        )
        .await
    }

    async fn fetch_authed(
        &self,
        db: &DB,
        username_override: Option<String>,
    ) -> error::Result<ApiAuthed> {
        fetch_api_authed(
            self.owner.clone(),
            self.email.clone(),
            &self.workspace_id,
            db,
            username_override,
        )
        .await
    }

    async fn disable_with_error(&self, db: &DB, error: String) -> () {
        if let Err(err) = sqlx::query!(
            "UPDATE capture_config SET error = $1, server_id = NULL, last_server_ping = NULL WHERE workspace_id = $2 AND path = $3 AND is_flow = $4 AND trigger_kind = 'websocket'",
            error,
            self.workspace_id,
            self.path,
            self.is_flow,
        )
        .execute(db).await {
            tracing::error!("Could not disable websocket capture {} ({}) with err {}, disabling because of error {}", self.path, self.workspace_id, err, error);
        }
    }

    fn get_trigger_path(&self) -> String {
        format!(
            "{}-{}",
            if self.is_flow { "flow" } else { "script" },
            self.path
        )
    }
}

enum WebsocketEnum {
    Trigger(WebsocketTrigger),
    Capture(CaptureConfigForWebsocket),
}

impl WebsocketEnum {
    async fn update_ping(&self, db: &DB, error: Option<&str>) -> Option<()> {
        match self {
            WebsocketEnum::Trigger(ws) => ws.update_ping(db, error).await,
            WebsocketEnum::Capture(capture) => capture.update_ping(db, error).await,
        }
    }

    async fn get_url_from_runnable(
        &self,
        path: &str,
        is_flow: bool,
        db: &DB,
    ) -> error::Result<String> {
        match self {
            WebsocketEnum::Trigger(ws) => ws.get_url_from_runnable(path, is_flow, db).await,
            WebsocketEnum::Capture(capture) => {
                capture.get_url_from_runnable(path, is_flow, db).await
            }
        }
    }

    async fn disable_with_error(&self, db: &DB, error: String) -> () {
        match self {
            WebsocketEnum::Trigger(ws) => ws.disable_with_error(db, error).await,
            WebsocketEnum::Capture(capture) => capture.disable_with_error(db, error).await,
        }
    }
}

async fn listen_to_websocket(
    ws: WebsocketEnum,
    db: DB,
    mut killpill_rx: tokio::sync::broadcast::Receiver<()>,
) -> () {
    if let None = ws.update_ping(&db, Some("Connecting")).await {
        return;
    }

    let url = match &ws {
        WebsocketEnum::Trigger(ws_trigger) => &ws_trigger.url,
        WebsocketEnum::Capture(capture) => &capture.trigger_config.url,
    };

    let filters: Vec<Filter> = match &ws {
        WebsocketEnum::Trigger(ws_trigger) => ws_trigger
            .filters
            .iter()
            .filter_map(|m| serde_json::from_str(m.get()).ok())
            .collect_vec(),
        WebsocketEnum::Capture(_) => vec![],
    };

    loop {
        let connect_url: Cow<str> = if url.starts_with("$") {
            if url.starts_with("$flow:") || url.starts_with("$script:") {
                let path = url.splitn(2, ':').nth(1).unwrap();
                tokio::select! {
                    biased;
                    _ = killpill_rx.recv() => {
                        return;
                    },
                    _ = loop_ping(&db, &ws, Some(
                        "Waiting on runnable to return websocket URL..."
                    )) => {
                        return;
                    },


                    url_result = ws.get_url_from_runnable(path, url.starts_with("$flow:"), &db) => match url_result {
                        Ok(url) => Cow::Owned(url),
                        Err(err) => {
                            ws.disable_with_error(&db, format!(
                                    "Error getting websocket URL from runnable after 5 tries: {:?}",
                                    err
                                ),
                            )
                            .await;
                            return;
                        }
                    },
                }
            } else {
                ws.disable_with_error(&db, format!("Invalid websocket runnable path: {}", url))
                    .await;
                return;
            }
        } else {
            Cow::Borrowed(url)
        };

        tokio::select! {
            biased;
            _ = killpill_rx.recv() => {
                return;
            },
            _ = loop_ping(&db, &ws, Some("Connecting...")) => {
                return;
            },
            connection = connect_async(connect_url.as_ref()) => {
                match connection {
                    Ok((ws_stream, _)) => {
                        tracing::info!("Listening to websocket {}", url);
                        if let None = ws.update_ping(&db, None).await {
                            return;
                        }
                        let (writer, mut reader) = ws_stream.split();
                        let mut last_ping = tokio::time::Instant::now();


                        // send initial messages
                        match &ws {
                            WebsocketEnum::Trigger(ws_trigger) => {
                                tokio::select! {
                                    biased;
                                    _ = killpill_rx.recv() => {
                                        return;
                                    },
                                    _ = loop_ping(&db, &ws, Some("Sending initial messages...")) => {
                                        return;
                                    },
                                    result = ws_trigger.send_initial_messages(writer, &db) => {
                                        if let Err(err) = result {
                                            ws_trigger.disable_with_error(&db, format!("Error sending initial messages: {:?}", err)).await;
                                            return
                                        } else {
                                            tracing::debug!("Initial messages sent successfully to websocket {}", url);
                                        }
                                    }
                                }
                            },
                            _ => {}
                        }

                        loop {
                            tokio::select! {
                                biased;
                                _ = killpill_rx.recv() => {
                                    return;
                                }
                                msg = reader.next() => {
                                    if let Some(msg) = msg {
                                        if last_ping.elapsed() > tokio::time::Duration::from_secs(5) {
                                            if let None = ws.update_ping(&db, None).await {
                                                return;
                                            }
                                            last_ping = tokio::time::Instant::now();
                                        }
                                        match msg {
                                            Ok(msg) => {
                                                match msg {
                                                    tokio_tungstenite::tungstenite::Message::Text(text) => {
                                                        tracing::debug!("Received text message from websocket {}: {}", url, text);
                                                        let mut should_handle = true;
                                                        for filter in &filters {
                                                            match filter {
                                                                Filter::JsonFilter(JsonFilter { key, value }) => {
                                                                    let mut deserializer = serde_json::Deserializer::from_str(text.as_str());
                                                                    should_handle = match is_value_superset(&mut deserializer, key, &value) {
                                                                        Ok(filter_match) => {
                                                                            filter_match
                                                                        },
                                                                        Err(err) => {
                                                                            tracing::warn!("Error deserializing filter for websocket {}: {:?}", url, err);
                                                                            false
                                                                        }
                                                                    };
                                                                }
                                                            }
                                                            if !should_handle {
                                                                break;
                                                            }
                                                        }
                                                        if should_handle {

                                                            let args = HashMap::from([("msg".to_string(), to_raw_value(&text))]);
                                                            let extra = Some(HashMap::from([(
                                                                "wm_trigger".to_string(),
                                                                to_raw_value(&serde_json::json!({"kind": "websocket", "websocket": { "url": url }})),
                                                            )]));

                                                            let args = PushArgsOwned { args, extra };
                                                            match &ws {
                                                                WebsocketEnum::Trigger(ws_trigger) => {
                                                                    ws_trigger.handle(&db, args).await;
                                                                },
                                                                WebsocketEnum::Capture(capture) => {
                                                                    capture.handle(&db, args).await;
                                                                }
                                                            }
                                                        }
                                                    },
                                                    a @ _ => {
                                                        tracing::debug!("Received non text-message from websocket {}: {:?}", url, a);
                                                    }
                                                }
                                            },
                                            Err(err) => {
                                                tracing::error!("Error reading from websocket {}: {:?}", url, err);
                                            }
                                        }
                                    } else {
                                        tracing::error!("Websocket {} closed", url);
                                        if let None = ws.update_ping(&db, Some("Websocket closed")).await {
                                            return;
                                        }
                                        tokio::time::sleep(tokio::time::Duration::from_secs(5)).await;
                                        break;
                                    }
                                },
                                _ = tokio::time::sleep(tokio::time::Duration::from_secs(5)) => {
                                    tracing::debug!("Sending ping to websocket {}", url);
                                    if let None = ws.update_ping(&db, None).await {
                                        return;
                                    }
                                    last_ping = tokio::time::Instant::now();
                                },
                            }
                        }
                    }
                    Err(err) => {
                        tracing::error!("Error connecting to websocket {}: {:?}", url, err);
                        if let None = ws.update_ping(&db, Some(err.to_string().as_str())).await {
                            return;
                        }
                        tokio::time::sleep(tokio::time::Duration::from_secs(5)).await;
                    }
                }
            }
        }
    }
}

async fn run_job(db: &DB, trigger: &WebsocketTrigger, args: PushArgsOwned) -> anyhow::Result<()> {
    let label_prefix = Some(format!("ws-{}-", trigger.path));

    let authed = fetch_api_authed(
        trigger.edited_by.clone(),
        trigger.email.clone(),
        &trigger.workspace_id,
        db,
        Some("anonymous".to_string()),
    )
    .await?;

    let user_db = UserDB::new(db.clone());

    let run_query = RunJobQuery::default();

    if trigger.is_flow {
        run_flow_by_path_inner(
            authed,
            db.clone(),
            user_db,
            trigger.workspace_id.clone(),
            StripPath(trigger.script_path.to_owned()),
            run_query,
            args,
            label_prefix,
        )
        .await?;
    } else {
        run_script_by_path_inner(
            authed,
            db.clone(),
            user_db,
            trigger.workspace_id.clone(),
            StripPath(trigger.script_path.to_owned()),
            run_query,
            args,
            label_prefix,
        )
        .await?;
    }

    Ok(())
}
