use std::{convert::Infallible, env, error::Error, net::SocketAddr, str::FromStr};

use base64::prelude::*;
use hmac::{digest::generic_array::functional::FunctionalSequence, Hmac, Mac};
use hyper::{
    service::{make_service_fn, service_fn}, Body, Client, HeaderMap, Method, Request, Response, Server, Uri
};
use serde::Deserialize;
use sha1::Sha1;

// STRUCTS


/// USER
struct User {
    name: String,
    roles: Vec<String>,
}


/// KeycloakToken
#[derive(Deserialize)]
struct KeycloakToken {
    pub access_token: String
}

/// DecodedJWT
#[derive(Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct DecodedJWT {
    #[serde(rename = "resource_access")]
    pub resource_access: Option<ResourceAccess>,
}

#[derive(Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct ResourceAccess {
    pub couchdb: Option<Couchdb>,
}

#[derive(Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct Couchdb {
    pub roles: Vec<String>,
}

// IMPLEMENTATIONS
impl User {
    fn couchdb_token(&self) -> String {
        let hmac_secret = env::var("COUCHDB_SECRET").unwrap();
        let mut hmac: Hmac<Sha1> =
            Mac::new_from_slice(hmac_secret.as_bytes()).unwrap();
        hmac.update(self.name.as_bytes());
        hmac.finalize()
            .into_bytes()
            .fold(String::new(), |acc, b| format!("{}{:02x}", acc, b))
    }

    pub fn set_headers(&self, headers: &mut HeaderMap) {
        headers.insert("X-Auth-CouchDB-UserName", self.name.parse().unwrap());
        headers.insert(
            "X-Auth-CouchDB-Roles",
            self.roles.join(",").parse().unwrap(),
        );
        headers.insert(
            "X-Auth-CouchDB-Token",
            self.couchdb_token().parse().unwrap()
        );
        headers.remove("authorization");
    }
}

impl KeycloakToken {
    fn from_json_text(text: &str) -> Result<KeycloakToken, serde_json::Error> {
        serde_json::from_str(text)
    }

    fn decode_token(&self) -> Result<DecodedJWT, Box<dyn Error>> {
        let data = self.access_token.split(".").skip(1).next();
        if data.is_none() {
            return Err("Invalid JWT".into());
        } 
        let decoded_text = String::from_utf8(BASE64_STANDARD_NO_PAD.decode(data.unwrap())?)?;
        Ok(serde_json::from_str(decoded_text.as_str())?)
    }
}

impl DecodedJWT {
    fn get_user(self, name: String) -> User {
        let roles = match self.resource_access {
            None => vec![],
            Some(ra) => {
                match ra.couchdb {
                    None => vec![],
                    Some(cdb) => cdb.roles
                }
            }
        };
        User {
            name,
            roles
        }
    }
}


// FUNCTIONS

async fn authenticate_keycloak(username: String, password: String) -> Result<User, (u16, String)> {
    let client_id = env::var("KEYCLOAK_CLIENT_ID").unwrap();
    let client_secret = env::var("KEYCLOAK_CLIENT_SECRET").unwrap();

    let mut headers = reqwest::header::HeaderMap::new();
    headers.insert("Content-Type", "application/x-www-form-urlencoded".parse().unwrap());

    let client = reqwest::Client::builder()
        .redirect(reqwest::redirect::Policy::none())
        .build()
        .unwrap();
    let res = client.post(env::var("KEYCLOAK_TOKEN_URL").unwrap())
        .headers(headers)
        .body(format!("grant_type=password&client_id={}&client_secret={}&username={}&password={}",client_id, client_secret, username, password))
        .send().await;

    if res.is_err() {
        return Err((400, "Cannot Connect to Keycloak".to_string()))
    }
    let res = res.unwrap();

    if res.status().as_u16() != 200 {
        return Err((res.status().as_u16(), res.status().to_string()))
    }

    let res_text = res.text().await;
    if let Err(e) = res_text {
        return Err((400, e.to_string()))
    }
    

    let keycloak_token = KeycloakToken::from_json_text(&res_text.unwrap());
    if let Err(e) = keycloak_token {
        return Err((400, e.to_string()))
    }

    let decoded_jwt = keycloak_token.unwrap().decode_token();
    if let Err(e) = decoded_jwt {
        return Err((400, e.to_string()));
    }

    Ok(decoded_jwt.unwrap().get_user(username))
}

fn extract_creds_from_request(req: &Request<Body>) -> Option<(String, String)> {
    let auth_value = req.headers().get("authorization")?;
    let b64_auth = auth_value.to_str().unwrap_or("").split(" ").skip(1).next()?;
    let decoded_auth = String::from_utf8(BASE64_STANDARD.decode(b64_auth.as_bytes()).unwrap_or(vec![])).unwrap();
    if let [username, password] = &decoded_auth.split(":").map(String::from).collect::<Vec<String>>()[..] {
        Some((username.to_string(), password.to_string()))
    } else {
        None
    }
}

fn forbidden_response() -> Response<Body> {
    Response::builder()
        .status(401)
        .header("www-authenticate", "Basic realm=\"CouchDB Proxy\"")
        .body(Body::empty())
        .unwrap()
}

fn error_response((status_code, msg) : (u16, String)) -> Response<Body> {
    Response::builder()
        .status(status_code)
        .body(Body::from(msg))
        .unwrap()
}

async fn handle(old_req: Request<Body>) -> Result<Response<Body>, hyper::Error> {
    let client = Client::new();
    
    // Build proxied Request
    let mut req = Request::from(old_req);
    let path = req.uri().path();
    *req.uri_mut() =
        Uri::from_str(format!("http://{}:{}{}", env::var("COUCHDB_HOST").unwrap(), env::var("COUCHDB_PORT").unwrap(), path.to_string()).as_str()).unwrap();

    if req.method() == Method::OPTIONS {
        return client.request(req).await
    }

    let creds = extract_creds_from_request(&req);
    if creds.is_none() {
        return Ok(forbidden_response());
    }
    let (username, password) = creds.unwrap();

    // Keycloak authentication
    let user = authenticate_keycloak(username, password).await;
    if let Err(e) = user {
        return Ok(error_response(e))
    }
    user.unwrap().set_headers(req.headers_mut());

    // Execute request
    match client.request(req).await {
        Ok(mut res) => {
            let res_headers = res.headers_mut();
            res_headers.insert("Via", "CouchDB Proxy".parse().unwrap());
            Ok(res)
        }
        Err(x) => Err(x),
    }
}

#[tokio::main]
async fn main() {
    let make_service = make_service_fn(|_| async { Ok::<_, Infallible>(service_fn(handle)) });
    let addr = SocketAddr::from(([0, 0, 0, 0], 8080));
    let server = Server::bind(&addr).serve(make_service);

    if let Err(e) = server.await {
        println!("Error: {}", e.message());
    }
}