use axum::http::HeaderMap;
use hmac::{Hmac, Mac};
use sha2::Sha256;
use std::fs;

type HmacSha256 = Hmac<Sha256>;

pub fn verify_signature(headers: HeaderMap, method: &str, path: &str) -> bool {
    let config = crate::s3_config::S3Config::load_default().unwrap_or_default();
    let mut config = config;
    config.merge_env();

    if !config.s3.require_auth {
        return true;
    }

    let auth_header = headers
        .get("Authorization")
        .and_then(|v| v.to_str().ok())
        .unwrap_or("");

    if !auth_header.starts_with("AWS4-HMAC-SHA256") {
        return false;
    }

    let credential = extract_credential(auth_header);
    if credential.is_none() {
        return false;
    }
    let credential = credential.unwrap();

    let secret_key = get_secret_key(&credential.access_key);
    if secret_key.is_none() {
        return false;
    }
    let secret_key = secret_key.unwrap();

    let x_amz_date = headers
        .get("X-Amz-Date")
        .and_then(|v| v.to_str().ok())
        .unwrap_or(&credential.date);

    let signed_headers = extract_signed_headers(auth_header);
    if signed_headers.is_none() {
        return false;
    }
    let signed_headers = signed_headers.unwrap();

    let payload_hash = get_payload_hash(&headers);

    let canonical_request = create_canonical_request(
        &headers,
        method,
        path,
        &signed_headers,
        &payload_hash,
    );

    let string_to_sign = create_string_to_sign(
        x_amz_date,
        &credential.region,
        &credential.service,
        &canonical_request,
    );

    let signing_key = calculate_signing_key(&secret_key, &credential.date, &credential.region, &credential.service);

    let calculated_signature = hmac_sha256_hex(&signing_key, &string_to_sign);

    let provided_signature = extract_signature(auth_header);
    if provided_signature.is_none() {
        return false;
    }

    calculated_signature == provided_signature.unwrap()
}

struct Credential {
    access_key: String,
    date: String,
    region: String,
    service: String,
}

fn extract_credential(auth_header: &str) -> Option<Credential> {
    let credential_part = auth_header
        .split(',')
        .find(|p| p.trim().starts_with("Credential="))?;

    let credential_str = credential_part.trim().strip_prefix("Credential=")?;
    let credential_parts: Vec<&str> = credential_str.split('/').collect();

    if credential_parts.len() < 5 {
        return None;
    }

    Some(Credential {
        access_key: credential_parts[0].to_string(),
        date: credential_parts[1].to_string(),
        region: credential_parts[2].to_string(),
        service: credential_parts[3].to_string(),
    })
}

fn extract_signed_headers(auth_header: &str) -> Option<Vec<String>> {
    let signed_headers_part = auth_header
        .split(',')
        .find(|p| p.trim().starts_with("SignedHeaders="))?;

    let signed_headers_str = signed_headers_part.trim().strip_prefix("SignedHeaders=")?;
    Some(signed_headers_str.split(';').map(|s| s.to_lowercase()).collect())
}

fn extract_signature(auth_header: &str) -> Option<String> {
    let signature_part = auth_header
        .split(',')
        .find(|p| p.trim().starts_with("Signature="))?;

    Some(signature_part.trim().strip_prefix("Signature=")?.to_string())
}

fn get_secret_key(access_key: &str) -> Option<String> {
    let s3_keys_path = "data/s3_keys.json";
    let s3_keys_json = fs::read_to_string(s3_keys_path).ok()?;

    #[derive(serde::Deserialize)]
    struct S3Key {
        access_key: String,
        secret_key: String,
    }

    let s3_keys: Vec<S3Key> = serde_json::from_str(&s3_keys_json).ok()?;
    s3_keys
        .iter()
        .find(|k| k.access_key == access_key)
        .map(|k| k.secret_key.clone())
}

fn get_payload_hash(headers: &HeaderMap) -> String {
    headers
        .get("X-Amz-Content-Sha256")
        .and_then(|v| v.to_str().ok())
        .map(|s| s.to_string())
        .unwrap_or_else(|| sha256_hex(""))
}

fn create_canonical_request(
    headers: &HeaderMap,
    method: &str,
    path: &str,
    signed_headers: &[String],
    payload_hash: &str,
) -> String {
    let canonical_uri = uri_encode(path, false);

    let canonical_query_string = build_canonical_query_string(headers);

    let canonical_headers = build_canonical_headers(headers, signed_headers);

    let signed_headers_str = signed_headers.join(";");

    format!(
        "{}\n{}\n{}\n{}\n{}\n{}",
        method,
        canonical_uri,
        canonical_query_string,
        canonical_headers,
        signed_headers_str,
        payload_hash
    )
}

fn uri_encode(input: &str, encode_slash: bool) -> String {
    input
        .chars()
        .map(|c| {
            if c.is_ascii_alphanumeric() || c == '-' || c == '_' || c == '.' || c == '~' {
                c.to_string()
            } else if c == '/' && !encode_slash {
                c.to_string()
            } else {
                format!("%{:02X}", c as u8)
            }
        })
        .collect()
}

fn build_canonical_query_string(_headers: &HeaderMap) -> String {
    // For S3, query string is typically empty for basic operations
    // This can be extended for presigned URLs
    String::new()
}

fn build_canonical_headers(headers: &HeaderMap, signed_headers: &[String]) -> String {
    signed_headers
        .iter()
        .map(|h| {
            let value = headers
                .get(h)
                .and_then(|v| v.to_str().ok())
                .unwrap_or("");
            format!("{}:{}\n", h, value.trim())
        })
        .collect()
}

fn create_string_to_sign(
    amz_date: &str,
    region: &str,
    service: &str,
    canonical_request: &str,
) -> String {
    let canonical_request_hash = sha256_hex(canonical_request);

    let date_stamp = &amz_date[..8];

    format!(
        "AWS4-HMAC-SHA256\n{}\n{}/{}/{}/aws4_request\n{}",
        amz_date,
        date_stamp,
        region,
        service,
        canonical_request_hash
    )
}

fn calculate_signing_key(secret_key: &str, date: &str, region: &str, service: &str) -> Vec<u8> {
    let k_secret = format!("AWS4{}", secret_key);
    let k_date = hmac_sha256(k_secret.as_bytes(), date);
    let k_region = hmac_sha256(&k_date, region);
    let k_service = hmac_sha256(&k_region, service);
    hmac_sha256(&k_service, "aws4_request")
}

fn hmac_sha256(key: &[u8], data: &str) -> Vec<u8> {
    let mut mac = HmacSha256::new_from_slice(key).expect("HMAC initialization failed");
    mac.update(data.as_bytes());
    mac.finalize().into_bytes().to_vec()
}

fn hmac_sha256_hex(key: &[u8], data: &str) -> String {
    let result = hmac_sha256(key, data);
    hex_encode(&result)
}

fn sha256_hex(data: &str) -> String {
    use sha2::Digest;
    let mut hasher = Sha256::new();
    hasher.update(data.as_bytes());
    let hash = hasher.finalize();
    hex_encode(&hash)
}

fn hex_encode(data: &[u8]) -> String {
    data.iter().map(|b| format!("{:02x}", b)).collect()
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_uri_encode() {
        assert_eq!(uri_encode("/bucket/key", false), "/bucket/key");
        assert_eq!(uri_encode("/bucket/key", true), "%2Fbucket%2Fkey");
        assert_eq!(uri_encode("test file.txt", false), "test%20file.txt");
    }

    #[test]
    fn test_sha256_hex() {
        let empty_hash = sha256_hex("");
        assert_eq!(
            empty_hash,
            "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855"
        );
    }

    #[test]
    fn test_calculate_signing_key() {
        let key = calculate_signing_key("secret", "20260621", "us-east-1", "s3");
        assert_eq!(key.len(), 32);
    }

    #[test]
    fn test_create_canonical_request() {
        let mut headers = HeaderMap::new();
        headers.insert("Host", "localhost:11438".parse().unwrap());
        let signed_headers = vec!["host".to_string()];
        
        let canonical = create_canonical_request(
            &headers,
            "GET",
            "/bucket/key",
            &signed_headers,
            "UNSIGNED-PAYLOAD",
        );
        
        assert!(canonical.contains("GET"));
        assert!(canonical.contains("host:localhost:11438"));
    }
}