use mumu::parser::interpreter::{Interpreter, DynamicFnInfo};
use mumu::parser::types::{
    FunctionValue, Value, InkIteratorHandle, InkIteratorKind, LavaStream,
};
use crate::register::credentials::SESSIONS;
use crate::functions::sigv4::{build_aws_sigv4_headers, hex_sha256_str};

use std::sync::{Arc, Mutex};
use std::thread;
use std::fs::File;
use std::io::Read;

use chrono::Utc;
use indexmap::IndexMap;
use reqwest::blocking::Client;
use reqwest::header::{HeaderValue, CONTENT_TYPE, USER_AGENT};
use xmltree::Element;

/* ────────────────────────────────────────────────────────────────────────── */
fn make_error(msg: impl Into<String>) -> Value {
    let mut m = IndexMap::new();
    m.insert("ok".to_string(), Value::Bool(false));
    m.insert("error".to_string(), Value::SingleString(msg.into()));
    Value::KeyedArray(m)
}
/* ────────────────────────────────────────────────────────────────────────── */

pub fn register_s3_fetch(interp: &mut Interpreter) {
    let f = Arc::new(Mutex::new(
        |interp: &mut Interpreter, args: Vec<Value>| -> Result<Value, String> {
            if interp.is_verbose() {
                eprintln!("[aws:s3:fetch] called with args: {:?}", args);
            }
            let res = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
                aws_s3_fetch_function(interp, args)
            }));
            match res {
                Ok(Ok(v)) => Ok(v),
                Ok(Err(msg)) => Ok(make_error(format!("aws:s3:fetch: {}", msg))),
                Err(e) => Ok(make_error(format!(
                    "aws:s3:fetch: fatal plugin panic: {:?}",
                    e
                ))),
            }
        },
    ));
    let info = DynamicFnInfo::new(f, false);
    interp.register_dynamic_function_ex("aws:s3:fetch", info);

    interp.set_variable(
        "aws:s3:fetch",
        Value::Function(Box::new(FunctionValue::Named("aws:s3:fetch".to_string()))),
    );
}

pub fn aws_s3_fetch_function(
    interp: &mut Interpreter,
    args: Vec<Value>,
) -> Result<Value, String> {
    if interp.is_verbose() {
        eprintln!("[aws:s3:fetch] entering function with args: {:?}", args);
    }
    /* ---------- Args ---------------------------------------------------- */
    if args.len() != 1 {
        return Err("expects a single keyed‑array argument".into());
    }
    let opts = match &args[0] {
        Value::KeyedArray(m) => m,
        _ => return Err("argument must be a keyed array".into()),
    };

    let cred_id = match opts.get("credentials") {
        Some(Value::SingleString(s)) => s,
        _ => return Err("missing 'credentials'".into()),
    };
    let bucket = match opts.get("bucket") {
        Some(Value::SingleString(s)) => s,
        _ => return Err("missing 'bucket'".into()),
    };
    let key = match opts.get("key") {
        Some(Value::SingleString(s)) => s,
        _ => return Err("missing 'key'".into()),
    };
    let path = opts
        .get("path")
        .and_then(|v| match v { Value::SingleString(s) => Some(s.clone()), _ => None });
    let folder = opts
        .get("folder")
        .and_then(|v| match v { Value::SingleString(s) => Some(s.clone()), _ => None });

    /* ---------- Session ------------------------------------------------- */
    let sessions = SESSIONS.lock().unwrap();
    let sess = sessions
        .get(cred_id)
        .ok_or_else(|| format!("no session with id {}", cred_id))?;

    /* ---------- Endpoint & SigV4 ---------------------------------------- */
    let host = format!("{}.s3.amazonaws.com", bucket);
    let endpoint_url = if let Some(ep) = &sess.endpoint {
        format!("{}/{}/{}", ep.trim_end_matches('/'), bucket, key)
    } else {
        format!("https://{}/{}", host, key)
    };

    let now = Utc::now();
    let amz_date = now.format("%Y%m%dT%H%M%SZ").to_string();
    let date_stamp = now.format("%Y%m%d").to_string();

    let (authz, mut headers) = build_aws_sigv4_headers(
        "GET",
        &format!("/{}", key),
        "",
        &hex_sha256_str(""),
        &host,
        &amz_date,
        &date_stamp,
        &sess.region,
        "s3",
        &sess.access_key,
        &sess.secret_key,
    );
    headers.insert(USER_AGENT, HeaderValue::from_static("aws-mumu/0.1"));
    headers.insert(CONTENT_TYPE, HeaderValue::from_static("application/octet-stream"));

    if interp.is_verbose() {
        eprintln!(
            "[aws:s3:fetch] Built SigV4 headers: {:?}, authz={}",
            headers, authz
        );
        eprintln!("[aws:s3:fetch] GET {}", endpoint_url);
    }

    let client = Client::builder()
        .build()
        .map_err(|e| format!("build client: {}", e))?;

    /* ---------- Direct download to file -------------------------------- */
    let save_path_opt = if let Some(p) = path {
        Some(p)
    } else if let Some(folder_path) = folder {
        use std::path::{Path, PathBuf};
        let mut p = PathBuf::from(folder_path);
        let filename = Path::new(&key)
            .file_name()
            .map(|s| s.to_string_lossy().to_string())
            .unwrap_or_else(|| "downloaded_file".into());
        p.push(filename);
        Some(p.to_string_lossy().to_string())
    } else {
        None
    };

    if let Some(save_path) = save_path_opt {
        let mut res = client
            .get(&endpoint_url)
            .headers(headers)
            .send()
            .map_err(|e| format!("HTTP error: {}", e))?;

        let status = res.status();
        if !status.is_success() {
            let body = res.text().unwrap_or_default();
            if let Ok(root) = Element::parse(body.as_bytes()) {
                let code = root
                    .get_child("Code")
                    .and_then(|c| c.get_text())
                    .unwrap_or_default();
                let msg = root
                    .get_child("Message")
                    .and_then(|c| c.get_text())
                    .unwrap_or_default();
                if !code.is_empty() || !msg.is_empty() {
                    return Err(format!("{}: {}", code, msg));
                }
            }
            return Err(format!("AWS error status {}", status)); // ← fixed
        }

        let mut out = File::create(&save_path)
            .map_err(|e| format!("create file: {}", e))?;
        res.copy_to(&mut out)
            .map_err(|e| format!("write file: {}", e))?;

        if interp.is_verbose() {
            eprintln!("[aws:s3:fetch] saved file to {}", save_path);
        }
        let mut ok = IndexMap::new();
        ok.insert("ok".into(), Value::Bool(true));
        ok.insert("path".into(), Value::SingleString(save_path));
        return Ok(Value::KeyedArray(ok));
    }

    /* ---------- Streaming mode ----------------------------------------- */
    let (tx, rx) = std::sync::mpsc::channel::<Vec<u8>>();
    let url_clone = endpoint_url.clone();
    let headers_clone = headers.clone();

    thread::spawn(move || {
        let client = Client::builder().build().unwrap();
        let mut res = match client.get(&url_clone).headers(headers_clone).send() {
            Ok(r) => r,
            Err(_) => return,
        };
        if !res.status().is_success() {
            return;
        }
        let mut buffer = [0u8; 8192];
        loop {
            match res.read(&mut buffer) {
                Ok(0) => break,
                Ok(n) => {
                    if tx.send(buffer[..n].to_vec()).is_err() {
                        break;
                    }
                }
                Err(_) => break,
            }
        }
    });

    let state = Arc::new(Mutex::new(S3FetchStreamState {
        rx,
        finished: false,
    }));

    let handle = InkIteratorHandle {
        kind: InkIteratorKind::Plugin(Arc::new(Mutex::new(Box::new(
            S3FetchStreamIter { state: state.clone() },
        ) as Box<dyn LavaStream>))),
    };

    if interp.is_verbose() {
        eprintln!("[aws:s3:fetch] returning InkIterator stream");
    }
    Ok(Value::InkIterator(handle))
}

/* ────────────────────────────────────────────────────────────────────────── */

struct S3FetchStreamState {
    rx: std::sync::mpsc::Receiver<Vec<u8>>,
    finished: bool,
}

struct S3FetchStreamIter {
    state: Arc<Mutex<S3FetchStreamState>>,
}

impl LavaStream for S3FetchStreamIter {
    fn next_value(&mut self) -> Result<Value, String> {
        let mut state = self.state.lock().unwrap();
        if state.finished {
            return Err("NO_MORE_DATA".into());
        }
        match state.rx.recv() {
            Ok(chunk) if !chunk.is_empty() => {
                let ints = chunk.into_iter().map(|b| b as i32).collect::<Vec<_>>();
                let mut m = IndexMap::new();
                m.insert("chunk".into(), Value::IntArray(ints));
                Ok(Value::KeyedArray(m))
            }
            _ => {
                state.finished = true;
                Err("NO_MORE_DATA".into())
            }
        }
    }
}
