=====src//client.rs=====
use log::{debug, error};
use std::env;
use reqwest::{Client};
use reqwest::header::{HeaderMap, HeaderValue, AUTHORIZATION};
use serde_json::{json, Value};
use std::time::Duration;
use crate::config::{FlowConfig, replace_with_env_var};


use serde::{Deserialize, Serialize};
use serde_json::Result;
use tokio::fs::File;
use tokio::io;
use tokio::io::AsyncReadExt;
use crate::client;
use serde_yaml::to_string as to_yaml;  // Add serde_yaml to your Cargo.toml if not already included


#[derive(Serialize, Deserialize, Debug)]
struct FluentCliOutput {
    pub(crate) text: String,
    pub(crate) question: String,
    #[serde(rename = "chatId")]
    pub(crate) chat_id: String,
    #[serde(rename = "chatMessageId")]
    chat_message_id: String,
    #[serde(rename = "sessionId")]
    pub(crate) session_id: String,
    #[serde(rename = "memoryType")]
    memory_type: Option<String>,
}

#[derive(Serialize, Deserialize, Debug)]
struct Question {
    question: String,
}


#[derive(Serialize, Deserialize)]
struct RequestPayload {
    question: String,
    overrideConfig: std::collections::HashMap<String, String>,
    uploads: Option<Vec<Upload>>,
}

#[derive(Serialize, Deserialize)]
struct Upload {
    data: String,
    r#type: String,
    name: String,
    mime: String,
}

#[derive(Debug)]
struct ResponseOutput {
    response_text: String,
    question: String,
    chat_id: String,
    session_id: String,
    memory_type: Option<String>,
    code_blocks: Option<Vec<String>>,  // Only populated if `--parse-code-output` is present
    pretty_text: Option<String>,       // Only populated if `--parse-code-output` is not present
}

use serde_json::Error as SerdeError;

pub async fn handle_response(response_body: &str, matches: &clap::ArgMatches) -> Result<()> {
    // Parse the response body, handle error properly here instead of unwrapping
    debug!("Response body: {}", response_body);
    let result = serde_json::from_str::<FluentCliOutput>(response_body);


    let response_text = match result {
        Ok(parsed_output) => {
            // If parsing is successful, use the parsed data
            debug!("{:?}", parsed_output);
            if let Some(directory) = matches.value_of("download-media") {
                let urls = extract_urls(response_body); // Assume extract_urls can handle any text
                download_media(urls, directory).await;
            }
            if matches.is_present("markdown-output") {
                pretty_format_markdown(&parsed_output.text); // Ensure text is obtained correctly

            } else if matches.is_present("parse-code-output") {
                let code_blocks = extract_code_blocks(&parsed_output.text);
                for block in code_blocks {
                    println!("{}", block);
                }
            } else if matches.is_present("full-output") {
                println!("{}", response_body);  // Output the text used, whether parsed or raw
            } else {
                println!("{}", parsed_output.text);  // Output the text used, whether parsed or raw, but only if the --markdown-output flag is not set").text;
            }
        },
        Err(e) => {
            // If there's an error parsing the JSON, print the error and the raw response body
            eprintln!("Raw Webhook Output");
            if let Some(directory) = matches.value_of("download-media") {
                let urls = extract_urls(response_body); // Assume extract_urls can handle any text
                download_media(urls, directory).await;
            }
            println!("{}", response_body);
            response_body.to_string();
        }
    };

    Ok(())
}


fn extract_urls(text: &str) -> Vec<String> {
    let url_regex = Regex::new(r"https?://[^\s]+").unwrap();
    url_regex.find_iter(text)
        .map(|mat| mat.as_str().to_string())
        .collect()
}

fn pretty_format_markdown(markdown_content: &str) {
    let skin = MadSkin::default(); // Assuming `termimad` is used
    skin.print_text(markdown_content); // Render to a string
}

fn extract_code_blocks(markdown_content: &str) -> Vec<String> {
    let re = Regex::new(r"```[\w]*\n([\s\S]*?)\n```").unwrap();
    re.captures_iter(markdown_content)
        .map(|cap| {
            cap[1].trim().to_string()  // Trim to remove leading/trailing whitespace
        })
        .collect()
}

pub fn parse_fluent_cli_output(json_data: &str) -> Result<FluentCliOutput> {
    let output: FluentCliOutput = serde_json::from_str(json_data)?;
    Ok(output)
}

use reqwest;
use tokio::io::AsyncWriteExt;
use chrono::Local;
use anyhow::{Context};

// Correct definition of the function returning a Result with a boxed dynamic error



async fn download_media(urls: Vec<String>, directory: &str) {
    let client = reqwest::Client::new();

    for url in urls {
        match client.get(&url).send().await {
            Ok(response) => {
                if response.status().is_success() {
                    match response.bytes().await {
                        Ok(content) => {
                            let path = Path::new(&url);
                            let filename = if let Some(name) = path.file_name() {
                                format!("{}-{}.{}", name.to_string_lossy(), Local::now().format("%Y%m%d%H%M%S"), path.extension().unwrap_or_default().to_string_lossy())
                            } else {
                                format!("download-{}.dat", Local::now().format("%Y%m%d%H%M%S"))
                            };
                            let filepath = Path::new(directory).join(filename);

                            match File::create(filepath).await {
                                Ok(mut file) => {
                                    if let Err(e) = file.write_all(&content).await {
                                        eprintln!("Failed to write to file: {}", e);
                                    }
                                },
                                Err(e) => eprintln!("Failed to create file: {}", e),
                            }
                        },
                        Err(e) => eprintln!("Failed to read bytes from response: {}", e),
                    }
                } else {
                    eprintln!("Failed to download {}: {}", url, response.status());
                }
            },
            Err(e) => eprintln!("Failed to send request: {}", e),
        }
    }
}

// Change the signature to accept a simple string for `question`

pub async fn send_request(flow: &FlowConfig,  payload: &Value) -> reqwest::Result<String> {
    let client = Client::new();

    // Dynamically fetch the bearer token from environment variables if it starts with "AMBER_"
    let bearer_token = if flow.bearer_token.starts_with("AMBER_") {
        env::var(&flow.bearer_token[6..]).unwrap_or_else(|_| flow.bearer_token.clone())
    } else {
        flow.bearer_token.clone()
    };
    debug!("Bearer token: {}", bearer_token);

    // Ensure override_config is up-to-date with environment variables
    let mut override_config = flow.override_config.clone();
    debug!("Override config before update: {:?}", override_config);
    replace_with_env_var(&mut override_config);
    debug!("Override config after update: {:?}", override_config);


    let url = format!("{}://{}:{}{}{}", flow.protocol, flow.hostname, flow.port, flow.request_path, flow.chat_id);
    debug!("URL: {}", url);
    debug!("Body: {}", payload);
    debug!("Headers: {:?}", bearer_token);
    // Send the request and await the response
    let response = client.post(&url)
        .header("Authorization", format!("Bearer {}", bearer_token))
        .json(payload)
        .send()
        .await?;

    debug!("Request URL: {}", url);
    debug!("Request bearer token: {}", bearer_token);
    debug!("Response: {:?}", response);

    response.text().await
}


pub(crate) fn build_request_payload(question: &str, context: Option<&str>) -> Value {
    // Construct the basic question
    let full_question = if let Some(ctx) = context {
        format!("{} {}", question, ctx)  // Concatenate question and context
    } else {
        question.to_string()  // Use question as is if no context
    };

    // Start building the payload with the question
    let mut payload = json!({
        "question": full_question,  // Use the potentially modified question
    });

    // Add the context to the payload if it exists
    if let Some(ctx) = context {
        payload.as_object_mut().unwrap().insert("context".to_string(), serde_json::Value::String(ctx.to_string()));
    }

    payload

}



use tokio::fs::File as TokioFile; // Alias to avoid confusion with std::fs::File
use tokio::io::{AsyncReadExt as TokioAsyncReadExt, Result as IoResult};
use base64::encode;
use std::collections::HashMap;
use std::io::ErrorKind;
use std::path::Path;
use pulldown_cmark::{Event, Parser, Tag};
use regex::Regex;
use serde::de::Error;
use termimad::{FmtText, MadSkin};
use termimad::minimad::once_cell::sync::Lazy;


pub(crate) async fn prepare_payload(flow: &FlowConfig, question: &str, file_path: Option<&str>, actual_final_context: Option<String>) -> IoResult<Value> {
    let mut override_config = flow.override_config.clone();
    // Ensure override_config is up-to-date with environment variables
    replace_with_env_var(&mut override_config);
    debug!("Override config after update: {:?}", override_config);

    debug!("File path: {:?}", file_path);
    debug!("Actual final context: {:?}", actual_final_context);

    let full_question = if let Some(ctx) = actual_final_context {
        format!("{} {}", question, ctx)  // Concatenate question and context
    } else {
        question.to_string()  // Use question as is if no context
    };
    // Assuming replace_with_env_var function exists and mutates override_config appropriately

    let mut body = json!({
        "question": full_question,
        "overrideConfig": override_config,
    });


    if let Some(path) = file_path {

        // Properly handle the file open result
        let mut file = TokioFile::open(path).await?;  // Correctly use .await and propagate errors with ?

        let mut buffer = Vec::new();
        // Use read_to_end on the file object directly
        TokioAsyncReadExt::read_to_end(&mut file, &mut buffer).await?;  // Correct usage with error propagation

        let encoded_image = encode(&buffer);  // Encode the buffer content to Base64
        let uploads = json!([{
            "data": format!("data:image/png;base64,{}", encoded_image),
            "type": "file",
            "name": path.rsplit('/').next().unwrap_or("unknown"),
            "mime": "image/png"
        }]);

        body.as_object_mut().unwrap().insert("uploads".to_string(), uploads);
    }

    Ok(body)
}

=====src//config.rs=====
use serde::{Deserialize, Serialize};
use serde_json::{Value, json};
use std::fs::File;
use std::io::Read;
use std::env;
use std::process::Command;
use std::str;
use std::error::Error;
use log::{debug, info};

#[derive(Debug, Deserialize, Serialize, Clone)]
pub struct FlowConfig {
    pub name: String,
    pub hostname: String,
    pub port: u16,
    pub chat_id: String,
    pub request_path: String,
    #[serde(rename = "sessionId")]
    pub session_id: String,
    pub bearer_token: String,
    #[serde(rename = "overrideConfig")]
    pub override_config: Value,
    pub timeout_ms: Option<u64>,
    pub protocol: String,
}



// Helper function to replace config strings starting with "AMBER_" with their env values
pub(crate) fn replace_with_env_var(value: &mut Value) {
    match value {
        Value::String(s) if s.starts_with("AMBER_") => {
            let env_key = &s[6..]; // Skip the "AMBER_" prefix to fetch the correct env var
            debug!("Attempting to replace: {}", s);
            debug!("Looking up environment variable: {}", env_key);
            match env::var(env_key) {
                Ok(env_value) => {
                    debug!("Environment value found: {}", env_value);
                    *s = env_value; // Successfully replace the string with the environment variable value.
                },
                Err(e) => {
                    debug!("Failed to find environment variable '{}': {}", env_key, e);
                    // Optionally, handle the error by logging or defaulting
                }
            }
        },
        Value::Object(map) => {
            for (k, v) in map.iter_mut() {
                debug!("Inspecting object key: {}", k);
                replace_with_env_var(v); // Recursively replace values in map
            }
        },
        Value::Array(arr) => {
            for item in arr.iter_mut() {
                replace_with_env_var(item); // Recursively replace values in array
            }
        },
        _ => {} // No action required for other types
    }
}



pub(crate) struct EnvVarGuard {
    keys: Vec<String>,
}

impl EnvVarGuard {
    pub fn new() -> Self {
        EnvVarGuard { keys: Vec::new() }
    }

    pub fn decrypt_amber_keys_for_flow(&mut self, flow: &mut FlowConfig) -> Result<(), Box<dyn Error>> {
        // Check each field in FlowConfig if it starts with "AMBER_"
        if flow.bearer_token.starts_with("AMBER_") {
            debug!("Decrypting bearer_token: {}", &flow.bearer_token);
            self.set_env_var_from_amber("bearer_token", &flow.bearer_token[6..])?;
        }
        if flow.session_id.starts_with("AMBER_") {
            debug!("Decrypting session_id: {}", &flow.session_id);
            self.set_env_var_from_amber("session_id", &flow.session_id[6..])?;
        }
        // Assume overrideConfig might also have Amber encrypted keys
        self.decrypt_amber_keys_in_value(&mut flow.override_config)?;
        debug!("Decrypted keys: {:?}", self.keys);
        Ok(())
    }

    fn decrypt_amber_keys_in_value(&mut self, value: &mut Value) -> Result<(), Box<dyn Error>> {
        debug!("Decrypting value: {:?}", value);
        if let Value::Object(obj) = value {
            debug!("Decrypting object: {:?}", obj);
            for (key, val) in obj.iter_mut() {
                debug!("Decrypting key: {}, value: {:?}", key, val);
                if let Some(str_val) = val.as_str() {
                    debug!("Decrypting str_val: {}", str_val);
                    if str_val.starts_with("AMBER_") {
                        debug!("Decrypting str_val: {}", str_val);
                        self.set_env_var_from_amber(key, &str_val[6..])?;
                        debug!("Decrypted key: {}, value: {:?}", key, val);
                    }
                } else {
                    self.decrypt_amber_keys_in_value(val)?;
                    debug!("Decrypted value: {:?}", val);
                }
            }
        }
        Ok(())
    }

    fn set_env_var_from_amber(&mut self, env_key: &str, amber_key: &str) -> Result<(), Box<dyn Error>> {
        let output = Command::new("amber")
            .args(&["print"])
            .output()?;
        debug!("Output of 'amber print {}': {}", amber_key, str::from_utf8(&output.stdout)?);

        if !output.status.success() {
            return Err(Box::new(std::io::Error::new(std::io::ErrorKind::Other, "Failed to decrypt Amber key")));
        }

        // This assumes the `amber print` outputs in the format: export AMBER_KEY_NAME="value"
        // This will find the line containing the amber_key and extract everything after the first '='
        let value = str::from_utf8(&output.stdout)?
            .lines()
            .find(|line| line.contains(amber_key))
            .ok_or("Amber key not found in the output")?
            .split_once('=')
            .ok_or("Invalid format of Amber output")?
            .1 // Take the part after the '='
            .trim() // Trim whitespace
            .trim_matches('"'); // Trim surrounding quotation marks if any

        // Use the amber_key as the environment variable name directly
        env::set_var(amber_key, value);  // Here, use amber_key instead of env_key
        debug!("Set environment variable {} with value {}", amber_key, value);
        self.keys.push(amber_key.to_owned());  // Store amber_key to track what has been set
        info!("Set environment variable {} with decrypted value.", amber_key);
        debug!("Keys: {:?}", self.keys);
        debug!("Environment variables: {:?}", env::vars());

        Ok(())
    }

}

impl Drop for EnvVarGuard {
    fn drop(&mut self) {
        for key in &self.keys {
            env::remove_var(key);
            info!("Environment variable {} has been unset.", key);
        }
    }
}




pub fn load_config() -> Result<Vec<FlowConfig>, Box<dyn Error>> {
    let config_path = env::var("FLUENT_CLI_CONFIG_PATH")
        .map_err(|_| "FLUENT_CLI_CONFIG_PATH environment variable is not set")?;
    let mut file = File::open(config_path)?;
    let mut contents = String::new();
    file.read_to_string(&mut contents)?;
    let configs: Vec<FlowConfig> = serde_json::from_str(&contents)?;

    Ok(configs)
}

pub fn generate_json_autocomplete_script() -> String {
    return format!(r#"
# Assuming FLUENT_CLI_CONFIG_PATH points to a JSON file containing configuration
autocomplete_flows() {{
    local current_word="${{COMP_WORDS[COMP_CWORD]}}"
    local flow_names=$(jq -r '.[].name' "$FLUENT_CLI_CONFIG_PATH")
    COMPREPLY=($(compgen -W "${{flow_names}}" -- "$current_word"))
}}
complete -F autocomplete_flows fluent_cli
"#)
}



pub fn generate_bash_autocomplete_script() -> String {

    return format!(r#"
# Assuming FLUENT_CLI_CONFIG_PATH points to a JSON file containing configuration
autocomplete_flows() {{
    local current_word="${{COMP_WORDS[COMP_CWORD]}}"
    local flow_names=$(jq -r '.[].name' "$FLUENT_CLI_CONFIG_PATH")
    COMPREPLY=($(compgen -W "${{flow_names}}" -- "$current_word"))
}}
complete -F autocomplete_flows fluent_cli
"#)

}


=====src//lib.rs=====
mod config;
mod client;=====src//main.rs=====
mod config;
mod client;

use ::config::Value;
use clap::{App, Arg, Command};
use tokio;

use log::{info, warn, error, debug};
use env_logger;
use tokio::fs::File;
use tokio::io::{self, AsyncReadExt};
use crate::client::handle_response;

use crate::config::{EnvVarGuard, generate_bash_autocomplete_script};
use anyhow::Result;


use colored::*; // Import the colored crate
use colored::control::*;

fn print_status(flowname: &str, new_question: &str) {
    eprintln!(
        "{}{}Fluent: {}\nProcessing: {}\n{}{}",
        "⫸⫸⫸⫸⫸".bright_yellow().bold(),
        "\n".normal(),
        flowname.bright_blue().bold(),
        new_question.bright_green(),
        "⫸⫸⫸⫸⫸".bright_yellow().bold(),
        "\n".normal()
    );
}

// use env_logger; // Uncomment this when you are using it to initialize logs
#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    env_logger::init();
    colored::control::set_override(true);

    let mut configs = config::load_config()?;

    let matches = Command::new("Fluent CLI")
        .version("0.1.0")
        .author("Nicholas Ferguson <nick@njf.io>")
        .about("Interacts with FlowiseAI workflows")
        .arg(Arg::new("flowname")
            .help("The flow name to invoke")
            .takes_value(true)
            .required_unless_present_any([
                "generate-autocomplete",
                "system-prompt-override-inline",
                "system-prompt-override-file",
                "additional-context-file"
            ]))
        .arg(Arg::new("request")
            .help("The request string to send")
            .takes_value(true)
            .required_unless_present_any([
                "generate-autocomplete",
                "system-prompt-override-inline",
                "system-prompt-override-file",
                "additional-context-file"
            ]))
        .arg(Arg::new("context")
            .short('c')  // Assigns a short flag
            .help("Optional context to include with the request")
            .takes_value(true))
        .arg(Arg::new("system-prompt-override-inline")
            .long("system-prompt-override-inline")
            .short('i')  // Assigns a short flag
            .help("Overrides the system message with an inline string")
            .takes_value(true))
        .arg(Arg::new("system-prompt-override-file")
            .long("system-prompt-override-file")
            .short('f')  // Assigns a short flag
            .help("Overrides the system message from a specified file")
            .takes_value(true))
        .arg(Arg::new("additional-context-file")
            .long("additional-context-file")
            .short('a')  // Assigns a short flag
            .help("Specifies a file from which additional request context is loaded")
            .takes_value(true))
        .arg(Arg::new("upload-image-path")
            .long("upload-image-path")
            .short('u')  // Assigns a short flag
            .value_name("FILE")
            .help("Sets the input file to use")
            .takes_value(true))
        .arg(Arg::new("generate-bash-autocomplete")
            .long("generate-autocomplete")
            .short('g')  // Assigns a short flag
            .help("Generates a bash autocomplete script")
            .takes_value(false))
        .arg(Arg::new("parse-code-output")
            .long("parse-code-output")
            .short('p')  // Assigns a short flag
            .help("Extracts and displays only the code blocks from the response")
            .takes_value(false))
        .arg(Arg::new("full-output")
            .long("full-output")
            .short('z')  // Assigns a short flag
            .help("Outputs all response data in JSON format")
            .takes_value(false))
        .arg(Arg::new("markdown-output")
            .long("markdown-output")
            .short('m')  // Assigns a short flag
            .help("Outputs the response to the terminal in stylized markdown. Do not use for pipelines")
            .takes_value(false))
        .arg(Arg::new("download-media")
            .long("download-media")
            .short('d')  // Assigns a short flag
            .help("Downloads all media files listed in the output to a specified directory")
            .takes_value(true)
            .value_name("DIRECTORY"))
        .get_matches();


    if matches.contains_id("generate-bash-autocomplete") {
        println!("{}", generate_bash_autocomplete_script());
        return Ok(());
    }

    let flowname = matches.value_of("flowname").unwrap();
    let flow = configs.iter_mut().find(|f| f.name == flowname).expect("Flow not found");
    let request = matches.value_of("request").unwrap();

    // Load context from stdin if not provided
    let context = matches.value_of("context");
    let mut additional_context = String::new();
    if context.is_none() && !atty::is(atty::Stream::Stdin) {
        tokio::io::stdin().read_to_string(&mut additional_context).await?;
    }
    debug!("Additional context: {:?}", additional_context);
    let final_context = context.or(if !additional_context.is_empty() { Some(&additional_context) } else { None });
    debug!("Context: {:?}", final_context);

    // Load override value from CLI if specified for system prompt override, file will always win
    let system_prompt_inline = matches.value_of("system-prompt-override-inline");
    let system_prompt_file = matches.value_of("system-prompt-override-file");
    // Load override value from file if specified
    let system_message_override = if let Some(file_path) = system_prompt_file {
        let mut file = File::open(file_path).await?; // Corrected async file opening
        let mut contents = String::new();
        file.read_to_string(&mut contents).await?; // Read file content asynchronously
        Some(contents)
    } else {
        system_prompt_inline.map(|s| s.to_string())
    };
    // Update the configuration with the override if it exists
    // Update the configuration based on what's present in the overrideConfig
    if let Some(override_value) = system_message_override {
        if let Some(obj) = flow.override_config.as_object_mut() {
            if obj.contains_key("systemMessage") {
                obj.insert("systemMessage".to_string(), serde_json::Value::String(override_value.to_string()));
            }
            if obj.contains_key("systemMessagePrompt") {
                obj.insert("systemMessagePrompt".to_string(), serde_json::Value::String(override_value.to_string()));
            }
        }
    }
    let file_path = matches.value_of("upload-image-path");



    // Determine the final context from various sources
    let file_context = matches.value_of("additional-context-file");
    let mut file_contents = String::new();
    if let Some(file_path) = file_context {
        let mut file = File::open(file_path).await?;
        file.read_to_string(&mut file_contents).await?;
    }


    // Combine file contents with other forms of context if necessary
    let actual_final_context = match (final_context, file_contents.is_empty()) {
        (Some(cli_context), false) => Some(format!("{} {}", cli_context, file_contents)),
        (None, false) => Some(file_contents),
        (Some(cli_context), true) => Some(cli_context.to_string()),
        (None, true) => None,
    };
    let actual_final_context_clone  = actual_final_context.clone();

    debug!("Actual Final context: {:?}", actual_final_context);
    let new_question = if let Some(ctx) = actual_final_context {
        format!("{} {}", request, ctx)  // Concatenate request and context
    } else {
        request.to_string()  // Use request as is if no context
    };


    // Decrypt the keys in the flow config
    let mut env_guard = EnvVarGuard::new();
    let env_guard_result = env_guard.decrypt_amber_keys_for_flow(flow)?;
    debug!("EnvGuard result: {:?}", env_guard_result);
    print_status(flowname, actual_final_context_clone.as_ref().unwrap_or(&new_question).as_str());
    let payload = crate::client::prepare_payload(&flow, request, file_path, actual_final_context_clone ).await?;
    let response = crate::client::send_request(&flow, &payload).await?;
    debug!("Handling Response");
    handle_response(response.as_str(), &matches).await?;
    Ok(())
}

