use crate::parser::{parse_authenticate_header, extract_next_page};
use base64::prelude::*;
use std::cell::RefCell;
use std::collections::HashMap;

pub trait Registries {
    fn get_tags(&self, registry: &str, image: &str) -> Option<Vec<String>>;
    fn get_digest(&self, registry: &str, image: &str, tag: &str) -> Option<String>;
    fn get_cache_ttl(&self, registry: &str) -> chrono::Duration;
}

#[derive(Debug)]
enum AuthMethod {
    Basic,
    Bearer { realm: String, service: String, scope: String },
}

impl AuthMethod {
    fn from_header(header: &ureq::http::HeaderValue) -> Option<Self> {
        let header_str = header.to_str().unwrap();
        let parse = parse_authenticate_header(header_str).unwrap();

        if &header_str[parse.scheme.clone()] == "Basic" {
            return Some(AuthMethod::Basic);
        }

        if &header_str[parse.scheme.clone()] == "Bearer" {
            return Some(AuthMethod::Bearer {
                realm: header_str[parse.realm.unwrap()].to_string(),
                scope: header_str[parse.scope.unwrap()].to_string(),
                service: header_str[parse.service.unwrap()].to_string(),
            });
        }

        return None;
    }
}

pub fn basic_auth(user: &str, pass: &str) -> String {
    return BASE64_STANDARD.encode(format!("{}:{}", user, pass));
}

pub enum Credentials {
    Anonymous,
    Basic(String),
}

pub struct Config {
    pub host: String,
    pub credentials: Credentials,
    pub cache_ttl: chrono::Duration,
}

enum AuthStage {
    Idle,
    Authorized(String),
}

struct RegistryState {
    config: Config,
    auth_stage: AuthStage,
}


impl RegistryState {
    fn add_to_request<T>(&self, req: ureq::RequestBuilder<T>) -> ureq::RequestBuilder<T> {
        if let AuthStage::Authorized(x) = &self.auth_stage {
            return req.header("Authorization", x);
        }

        return req;
    }

    fn authenticate(&mut self, response: &ureq::http::Response<ureq::Body>) -> Result<(), String> {
        match self.auth_stage {
            AuthStage::Authorized(_) =>
                // The token must have expired
                self.auth_stage = AuthStage::Idle,
            AuthStage::Idle => {}
        }

        let auth_header = response.headers().get("www-authenticate").unwrap();
        match AuthMethod::from_header(auth_header) {
            Some(AuthMethod::Basic) => {
                match &self.config.credentials {
                    Credentials::Basic(creds) => {
                        self.auth_stage = AuthStage::Authorized(format!("Basic {}", creds));
                        return Ok(());
                    }
                    Credentials::Anonymous => {
                        return Err(format!(
                            "Server {} requires Basic auth but configured as anonymous",
                            self.config.host
                        ));
                    }
                }
            }
            Some(AuthMethod::Bearer { realm, scope, service }) => {
                let url = format!("{}?service={}&scope={}", realm, service, scope);

                let mut request = ureq::get(url)
                    .config()
                    .http_status_as_error(false)
                    .build();

                if let Credentials::Basic(creds) = &self.config.credentials {
                    request = request.header("Authorization", format!("Basic {}", creds));
                }

                let body = request.call();

                let body = body
                    .map_err(|x| {
                        format!(
                            "Server {} authentication request failed: {}",
                            self.config.host,
                            x.to_string()
                        )
                    })?
                    .body_mut()
                    .read_to_string()
                    .expect(
                        format!(
                            "Server {} responded with something non-string like",
                            self.config.host
                        )
                        .as_str(),
                    );

                let body = body.parse::<tinyjson::JsonValue>().unwrap();

                let token = body["token"].get::<String>().unwrap();
                self.auth_stage = AuthStage::Authorized(format!("Bearer {}", token));
                return Ok(());
            }
            None => {
                return Err(format!(
                    "Server {} provided us with a challenge, but we didn't understand it",
                    self.config.host
                ))
            }
        }
    }
}

#[derive(Debug)]
enum RegistryError {
    NotFound,
    Other(String),
}

pub struct HttpRegistry {
    auth: RefCell<HashMap<String, RegistryState>>,
}

impl HttpRegistry {
    pub fn new(infos: Vec<Config>) -> Self {
        let mut states = HashMap::new();
        for info in infos {
            states.insert(info.host.clone(), RegistryState {
                config: info,
                auth_stage: AuthStage::Idle,
            });
        }
        return HttpRegistry { auth: RefCell::new(states) };
    }

    fn perform_request(
        &self,
        registry: &str,
        url: &str,
        accept: &'static str,
    ) -> Result<ureq::http::Response<ureq::Body>, RegistryError> {
        let mut auth = self.auth.borrow_mut();
        let mut state = auth.get_mut(registry);

        let url = format!("https://{}{}", registry, &url);

        let mut authentication_retry = false;

        loop {
            let mut request = ureq::get(&url)
                .header("Accept", accept)
                .config()
                .http_status_as_error(false)
                .build();

            if let Some(ref state) = state {
                request = state.add_to_request(request);
            }

            let response = {
                let _timer = crate::metrics::get()
                    .http_request_duration
                    .with_label_values(&[registry])
                    .start_timer();
                request.call().unwrap()
            };
            crate::metrics::get()
                .http_request
                .with_label_values(&[registry])
                .inc();

            if response.status() == 401 {
                if !authentication_retry {
                    if let Some(ref mut state) = state {
                        match state.authenticate(&response) {
                            Ok(()) => {
                                authentication_retry = true;
                                continue;
                            }
                            Err(e) => {
                                return Err(RegistryError::Other(e));
                            }
                        }
                    } else {
                        return Err(RegistryError::Other(format!(
                            "Server {} returned 401 but we have no credentials",
                            registry
                        )));
                    }
                } else {
                    return Err(RegistryError::Other(format!("Authentication failed")));
                }
            }

            if response.status() == 429 {
                crate::metrics::get().state.set(crate::metrics::STATE_THROTTLED);
                println!("Too many requests");
                std::thread::sleep(std::time::Duration::from_secs(8));
                crate::metrics::get().state.set(crate::metrics::STATE_ACTIVE);
                continue;
            }

            if response.status() == 404 {
                return Err(RegistryError::NotFound);
            }

            if response.status() != 200 {
                return Err(RegistryError::Other(format!(
                    "Unexpected status code: {}",
                    response.status()
                )));
            }

            return Ok(response);
        }
    }
}

impl Registries for HttpRegistry {
    fn get_tags(&self, registry: &str, image: &str) -> Option<Vec<String>> {
        let mut url = format!("/v2/{}/tags/list", image);
        let mut tags = vec![];

        loop {
            let mut response = match self.perform_request(
                registry,
                &url,
                "application/vnd.oci.image.index.v1+json",
            ) {
                Ok(r) => r,
                Err(RegistryError::NotFound) => return None,
                Err(RegistryError::Other(msg)) => {
                    println!("Error fetching tags: {}", msg);
                    return None;
                }
            };

            let content_type = response.headers()["Content-Type"].to_str().unwrap();
            if !content_type.starts_with("application/json") {
                println!("Unexpected Content-Type: {}", content_type);
                return None;
            }

            let body = response.body_mut().read_to_string().unwrap();
            let body = body.parse::<tinyjson::JsonValue>().unwrap();

            for it in body["tags"].get::<Vec<tinyjson::JsonValue>>().unwrap().iter() {
                tags.push(it.get::<String>().unwrap().clone());
            }

            if let Some(link_header) = response.headers().get("link") {
                let link_str = link_header.to_str().unwrap();
                let link = extract_next_page(&link_str).unwrap();
                url = link_str[link.next_uri.unwrap().clone()].to_string();
            } else {
                break;
            }
        }

        return Some(tags);
    }

    fn get_digest(&self, registry: &str, image: &str, tag: &str) -> Option<String> {
        let url = format!("/v2/{}/manifests/{}", image, tag);
        let response = match self.perform_request(
            registry,
            &url,
            "application/vnd.oci.image.manifest.v1+json,application/vnd.oci.image.index.v1+json",
        ) {
            Ok(r) => r,
            Err(RegistryError::NotFound) => return None,
            Err(RegistryError::Other(msg)) => {
                println!("Error fetching digest: {}", msg);
                return None;
            }
        };

        let digest = response.headers()["docker-content-digest"]
            .to_str()
            .unwrap()
            .to_string();
        return Some(digest);
    }

    fn get_cache_ttl(&self, registry: &str) -> chrono::Duration {
        let auth = self.auth.borrow();
        if let Some(state) = auth.get(registry) {
            return state.config.cache_ttl;
        }
        return chrono::Duration::minutes(1440);
    }
}

#[cfg(test)]
struct StubTag {
    registry: String,
    image: String,
    tag: String,
    digest: String,
}

#[cfg(test)]
pub struct StubRegistry {
    entries: RefCell<Vec<StubTag>>,
}

#[cfg(test)]
impl Default for StubRegistry {
    fn default() -> Self {
        return StubRegistry {
            entries: RefCell::new(vec![]),
        };
    }
}

#[cfg(test)]
impl StubRegistry {
    pub fn add_tag(&self, registry: &str, image: &str, tag: &str, digest: &str) {
        self.entries.borrow_mut().push(StubTag {
            registry: registry.to_string(),
            image: image.to_string(),
            tag: tag.to_string(),
            digest: digest.to_string(),
        });
    }

    pub fn set_digest(&self, registry: &str, image: &str, tag: &str, digest: &str) {
        for entry in self.entries.borrow_mut().iter_mut() {
            if entry.registry == registry && entry.image == image && entry.tag == tag {
                entry.digest = digest.to_string();
                return;
            }
        }
    }
}

#[cfg(test)]
impl Registries for StubRegistry {
    fn get_tags(&self, registry: &str, image: &str) -> Option<Vec<String>> {
        let mut tags = vec![];
        let mut found = false;

        for entry in self.entries.borrow().iter() {
            if entry.registry == registry && entry.image == image {
                tags.push(entry.tag.clone());
                found = true;
            }
        }

        if !found {
            return None;
        }

        return Some(tags);
    }

    fn get_digest(&self, registry: &str, image: &str, tag: &str) -> Option<String> {
        for entry in self.entries.borrow().iter() {
            if entry.registry == registry && entry.image == image && entry.tag == tag {
                return Some(entry.digest.clone());
            }
        }
        return None;
    }

    fn get_cache_ttl(&self, _registry: &str) -> chrono::Duration {
        return chrono::Duration::minutes(1440);
    }
}

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

    fn test_get_tags_returns_none_for_unknown(reg: &dyn Registries) {
        assert_eq!(reg.get_tags("docker.io", "unknown"), None);
    }

    fn test_get_tags_returns_all(reg: &dyn Registries, stub: &StubRegistry) {
        stub.add_tag("docker.io", "nginx", "1.21", "sha256:abc");
        stub.add_tag("docker.io", "nginx", "1.20", "sha256:def");
        stub.add_tag("docker.io", "nginx", "latest", "sha256:ghi");
        let mut tags = reg.get_tags("docker.io", "nginx").unwrap();

        tags.sort();
        assert_eq!(tags, vec!["1.20", "1.21", "latest"]);
    }

    fn test_get_digest_returns_none_for_unknown(reg: &dyn Registries) {
        assert_eq!(reg.get_digest("docker.io", "nginx", "1.21"), None);
    }

    fn test_get_digest_returns_digest(reg: &dyn Registries, stub: &StubRegistry) {
        stub.add_tag("docker.io", "nginx", "1.21", "sha256:abc123");

        assert_eq!(reg.get_digest("docker.io", "nginx", "1.21"), Some("sha256:abc123".to_string()));
    }

    fn test_set_digest_updates_existing(reg: &dyn Registries, stub: &StubRegistry) {
        stub.add_tag("docker.io", "nginx", "1.21", "sha256:old");
        stub.set_digest("docker.io", "nginx", "1.21", "sha256:new");

        assert_eq!(reg.get_digest("docker.io", "nginx", "1.21"), Some("sha256:new".to_string()));
    }

    #[test]
    fn conformance_stub_get_tags_unknown() {
        let stub = StubRegistry::default();
        test_get_tags_returns_none_for_unknown(&stub);
    }

    #[test]
    fn conformance_stub_get_tags_all() {
        let stub = StubRegistry::default();
        test_get_tags_returns_all(&stub, &stub);
    }

    #[test]
    fn conformance_stub_get_digest_unknown() {
        let stub = StubRegistry::default();
        test_get_digest_returns_none_for_unknown(&stub);
    }

    #[test]
    fn conformance_stub_get_digest() {
        let stub = StubRegistry::default();
        test_get_digest_returns_digest(&stub, &stub);
    }

    #[test]
    fn conformance_stub_set_digest() {
        let stub = StubRegistry::default();
        test_set_digest_updates_existing(&stub, &stub);
    }

    #[cfg(feature = "integration")]
    fn test_http_get_tags(reg: &dyn Registries) {
        let tags = reg.get_tags("registry.hub.docker.com", "library/nginx");
        assert!(tags.is_some());
        let tags = tags.unwrap();
        assert!(tags.contains(&"1.21".to_string()));
    }

    #[cfg(feature = "integration")]
    fn test_http_get_digest(reg: &dyn Registries) {
        let digest = reg.get_digest("registry.hub.docker.com", "library/nginx", "1.21");
        assert!(digest.is_some());
        assert!(digest.unwrap().starts_with("sha256:"));
    }

    #[test]
    #[cfg(feature = "integration")]
    fn conformance_http_get_tags() {
        crate::metrics::init();
        let reg = HttpRegistry::new(vec![AuthInfo {
            host: "registry.hub.docker.com".to_string(),
            credentials: Credentials::Anonymous,
        }]);
        test_http_get_tags(&reg);
    }

    #[test]
    #[cfg(feature = "integration")]
    fn conformance_http_get_digest() {
        crate::metrics::init();
        let reg = HttpRegistry::new(vec![AuthInfo {
            host: "registry.hub.docker.com".to_string(),
            credentials: Credentials::Anonymous,
        }]);
        test_http_get_digest(&reg);
    }
}