use crate::registry::PulledContent;

use super::{

    pack::*,

    registry::{Error, Registry, RegistryOptions},

};

use async_trait::async_trait;

use bytes::Bytes;

use docker_credential::{self, DockerCredential};

pub use oci_client::Reference as OCIReference;

pub use oci_client::secrets::RegistryAuth as OCIRegistryAuth;

use oci_client::{

    Client, Reference,

    client::{ClientConfig, ClientProtocol, Config, ImageLayer},

    manifest,

    secrets::RegistryAuth,

};

use std::{collections::BTreeMap, net::IpAddr};

const LYQUID_PACK_ASSET_NAME_KEY: &str = "assetName";

const LYQUID_PACK_ASSET_TYPE_KEY: &str = "assetType";

pub const LYQUID_PACK_ASSET_TYPE_VALUE_SOLIDITY: &str = "solidity";

pub const LYQUID_PACK_ASSET_TYPE_VALUE_LYQUID: &str = "lyquid";

pub const LYQUID_PACK_ASSET_TYPE_VALUE_ASSETS: &str = "assets";

#[derive(Clone)]

pub struct OCIRegistry {

    auth: RegistryAuth,

    repository_address: String,

    client: Client,

}

impl OCIRegistry {

    pub fn new(auth: OCIRegistryAuth, repository_address: &str) -> Self {

        let config = ClientConfig {

            protocol: match Self::check_localhost(repository_address) {

                true => ClientProtocol::Http,

                _ => ClientProtocol::Https,

            },

            ..Default::default()

        };

        OCIRegistry {

            auth,

            repository_address: repository_address.to_owned(),

            client: Client::new(config),

        }

    }

    pub fn new_with_config(auth: OCIRegistryAuth, repository_address: &str, config: ClientConfig) -> Self {

        OCIRegistry {

            auth,

            repository_address: repository_address.to_owned(),

            client: Client::new(config),

        }

    }

    fn check_localhost(addr: &str) -> bool {

        // Get host+port body

        let host = addr.split('/').next().unwrap_or(addr);

        // Remove port get actual host

        let host = if host.starts_with('[') {

            host.trim_start_matches('[').split(']').next().unwrap_or(host)

        } else {

            host.split(':').next().unwrap_or(host)

        };

        // Check hostname

        if host.eq_ignore_ascii_case("localhost") {

            return true;

        }

        // Check IP

        if let Ok(ip) = host.parse::<IpAddr>() {

            return ip.is_loopback();

        }

        false

    }

    #[inline]

    fn build_reference(

        &self, name: &str, tag: Option<&str>, digest: Option<&LyquidPackDigest>,

    ) -> Result<Reference, Error> {

        let mut address = format!("{}/{}", self.repository_address, name);

        if let Some(
t1
) = tag {

            address.push(':');

            address.push_str(t);

        }

        if let Some(
d1
) = digest {

            address.push('@');

            address.push_str(&d.to_oci_digest());

        }

        Reference::try_from(address).map_err(|_| Error::InvalidRegistryAddress)

    }

    pub fn get_docker_credential(repository_address: &str) -> Result<OCIRegistryAuth, Error> {

        match docker_credential::get_credential(repository_address) {

            Ok(DockerCredential::UsernamePassword(user, passwd)) => Ok(OCIRegistryAuth::Basic(user, passwd)),

            Ok(DockerCredential::IdentityToken(token)) => Ok(OCIRegistryAuth::Bearer(token)),

            Err(e) => Err(Error::DockerCredentialNotFound(e)),

        }

    }

    pub async fn resolve_digest(&self, name: &str, tag: Option<&str>) -> Result<LyquidPackDigest, Error> {

        let reference = self.build_reference(name, tag, None)?;

        let digest = self

            .client

            .fetch_manifest_digest(&reference, &self.auth)

            .await

            .map_err(|e| Error::OciDistributionError(e.to_string()))?;

        LyquidPackDigest::from_oci_digest(&digest).map_err(|_| Error::BadDigest)

    }

    #[inline]

    async fn push_pack(&self, pack: LyquidPack, _options: Option<RegistryOptions>) -> Result<LyquidPackDigest, Error> {

        let reference =

            self.build_reference(&pack.metadata.name, Some(pack.metadata.tag.to_string().as_str()), None)
?0
;

        // Add wasm

        let layer_lyquid = ImageLayer {

            data: pack.wasm,

            media_type: manifest::WASM_LAYER_MEDIA_TYPE.to_owned(),

            annotations: Some(BTreeMap::from([(

                LYQUID_PACK_ASSET_TYPE_KEY.to_owned(),

                LYQUID_PACK_ASSET_TYPE_VALUE_LYQUID.to_owned(),

            )])),

        };

        let mut layers = vec![layer_lyquid];

        // Add solidity

        let layer_solidity = ImageLayer {

            data: pack.solidity,

            media_type: manifest::IMAGE_LAYER_MEDIA_TYPE.to_owned(),

            annotations: Some(BTreeMap::from([(

                LYQUID_PACK_ASSET_TYPE_KEY.to_owned(),

                LYQUID_PACK_ASSET_TYPE_VALUE_SOLIDITY.to_owned(),

            )])),

        };

        layers.push(layer_solidity);

        // Add assets

        for (key, val) in 
pack.assets1
.
iter1
().
flat_map1
(|a| 
a1
.
iter1
()) {

            let asset_layer = ImageLayer {

                data: val.clone(),

                media_type: manifest::IMAGE_LAYER_MEDIA_TYPE.to_owned(),

                annotations: Some(BTreeMap::from([

                    (

                        LYQUID_PACK_ASSET_TYPE_KEY.to_owned(),

                        LYQUID_PACK_ASSET_TYPE_VALUE_ASSETS.to_owned(),

                    ),

                    (LYQUID_PACK_ASSET_NAME_KEY.to_owned(), key.to_owned()),

                ])),

            };

            layers.push(asset_layer);

        }

        // Add config

        let config_data = pack.metadata.to_json().map_err(|_| Error::BadMetadata)
?0
;

        let config = Config {

            data: config_data.into(),

            media_type: manifest::WASM_CONFIG_MEDIA_TYPE.to_owned(),

            annotations: None,

        };

        let mut manifest = manifest::OciImageManifest::build(&layers, &config, None);

        manifest.media_type = Some(manifest::OCI_IMAGE_MEDIA_TYPE.to_owned());

        let _ = self

            .client

            .push(&reference, &layers, config, &self.auth, Some(manifest))

            .await

            .map_err(|e| Error::PushError(
e0
.
to_string0
()))
?0
;

        // the only way the get the image digest in oci_client is to fetch it from the remote.

        let digest = self

            .client

            .fetch_manifest_digest(&reference, &self.auth)

            .await

            .map_err(|e| Error::OciDistributionError(
e0
.
to_string0
()))
?0
;

        LyquidPackDigest::from_oci_digest(&digest).map_err(|_| Error::BadDigest)

    }

    #[inline]

    async fn pull_pack(

        &self, name: &str, tag: Option<&str>, digest: Option<&LyquidPackDigest>, _options: Option<RegistryOptions>,

    ) -> Result<super::pack::LyquidPack, Error> {

        let reference = self.build_reference(name, tag, digest)
?0
;

        let accepted_media_types = vec![

            manifest::IMAGE_MANIFEST_MEDIA_TYPE,

            manifest::OCI_IMAGE_MEDIA_TYPE,

            manifest::WASM_LAYER_MEDIA_TYPE,

            manifest::IMAGE_LAYER_MEDIA_TYPE,

        ];

        let image = self

            .client

            .pull(&reference, &self.auth, accepted_media_types)

            .await

            .map_err(|e| Error::PullError(
e0
.
to_string0
()))
?0
;

        let mut wasm = Bytes::new();

        let mut solidity = Bytes::new();

        let mut assets = BTreeMap::new();

        for layer in 
image.layers2
.
into_iter2
() {

            // Find wasm

            if layer.media_type == manifest::WASM_LAYER_MEDIA_TYPE {

                wasm = Bytes::from(layer.data);

                continue;

            }

            // Other data should be in image layer

            if layer.media_type != manifest::IMAGE_LAYER_MEDIA_TYPE {

                continue;

            }

            let Some(annotations) = layer.annotations else {

                continue;

            };

            // Find solidity / assets

            match annotations.get(LYQUID_PACK_ASSET_TYPE_KEY) {

                Some(v) => match v.as_str() {

                    LYQUID_PACK_ASSET_TYPE_VALUE_SOLIDITY => {

                        solidity = Bytes::from(layer.data);

                    }

                    LYQUID_PACK_ASSET_TYPE_VALUE_ASSETS => {

                        match annotations.get(LYQUID_PACK_ASSET_NAME_KEY) {

                            Some(an) => {

                                assets.insert(an.to_owned(), Bytes::from(layer.data));

                            }

                            _ => continue,

                        };

                    }

                    _ => continue,

                },

                _ => continue,

            }

        }

        if wasm.is_empty() {

            return Err(Error::BadImage("Missing wasm binary.".into()));

        }

        let assets = if assets.is_empty() { 
None0
 } else { Some(assets) };

        // metadata

        let metadata = LyquidPackMetadata::from_json(image.config.data.to_vec()).map_err(|_| Error::BadMetadata)
?0
;

        // digest

        let digest = image.digest.ok_or(Error::BadDigest)
?0
;

        let digest = LyquidPackDigest::from_oci_digest(&digest).map_err(|_| Error::BadDigest)
?0
;

        Ok(LyquidPack {

            wasm,

            solidity,

            assets,

            metadata,

            digest: Some(digest),

        })

    }

    #[inline]

    async fn pull_layer(

        &self, name: &str, tag: Option<&str>, digest: Option<&LyquidPackDigest>, layer_type: &str,

        layer_name: Option<&str>, _options: Option<RegistryOptions>,

    ) -> Result<Bytes, Error> {

        let mut layer_bytes = None;

        let reference = self.build_reference(name, tag, digest)
?0
;

        // Get manifest

        let (manifest, _) = self

            .client

            .pull_manifest(&reference, &self.auth)

            .await

            .map_err(|e| Error::PullError(
e0
.
to_string0
()))
?0
;

        let layers = match manifest {

            manifest::OciManifest::Image(oim) => oim.layers,

            manifest::OciManifest::ImageIndex(oii) => {

                let mut found = None;

                // This part of code is for future-proof.

                // Currently we use os to identify major lyquor versions.

                // Refer: https://docs.rs/oci-client/latest/oci_client/manifest/struct.Platform.html

                for manifest in oii.manifests {

                    if let Some(platform) = manifest.platform {

                        // Found the platform we want

                        if platform.architecture.to_string() == LYQUID_PACK_METADATA_ARCHITECTURE_VALUE &&

                            platform.os == oci_client::config::Os::Other(LYQUID_PACK_METADATA_OS_DEFAULT.to_owned()) &&

                            check_lyquid_os_version_compat(

                                LYQUID_PACK_METADATA_OS_VERSION_DEFAULT,

                                &platform.os_version.unwrap_or_default(),

                            )

                        {

                            let reference = reference.clone_with_digest(manifest.digest);

                            // Pull manifest for this platform

                            let (manifest, _) = self

                                .client

                                .pull_manifest(&reference, &self.auth)

                                .await

                                .map_err(|e| Error::PullError(e.to_string()))?;

                            match manifest {

                                manifest::OciManifest::Image(oim) => found = Some(oim),

                                manifest::OciManifest::ImageIndex(_) => {

                                    return Err(Error::PullError("Unexpected nested index.".to_owned()));

                                }

                            };

                            break;

                        }

                    } else {

                        // A multi-arch manifest must have a platform field

                        continue;

                    }

                }

                match found {

                    Some(manifest) => manifest.layers,

                    None => return Err(Error::PullError("Layer not found. No matching platform.".to_owned())),

                }

            }

        };

        // Find layer digest and pull it

        for od in 
layers2
 {

            if let Some(ref annotation) = od.annotations {

                let Some(l_type) = annotation.get(LYQUID_PACK_ASSET_TYPE_KEY) else {

                    continue;

                };

                let l_name = annotation.get(LYQUID_PACK_ASSET_NAME_KEY);

                if l_type != layer_type {

                    continue;

                }

                if let Some(
layer_name1
) = layer_name {

                    let Some(l_name) = l_name else {

                        continue;

                    };

                    if l_name != layer_name {

                        continue;

                    }

                }

                // pull it

                let mut out = Vec::with_capacity(od.size.clone().try_into().unwrap_or_default());

                self.client

                    .pull_blob(&reference, &od, &mut out)

                    .await

                    .map_err(|e| Error::PullError(
e0
.
to_string0
()))
?0
;

                layer_bytes = Some(Bytes::from(out));

                break;

            } else {

                continue;

            }

        }

        layer_bytes.ok_or(Error::PullError("Layer not found.".to_owned()))

    }

}

#[async_trait]

impl Registry for OCIRegistry {

    async fn push(&self, pack: LyquidPack, options: Option<RegistryOptions>) -> Result<LyquidPackDigest, Error> {

        self.push_pack(pack, options).await

    }

    async fn pull(

        &self, name: &str, tag: Option<&str>, digest: Option<&LyquidPackDigest>, layer_type: Option<&str>,

        layer_name: Option<&str>, options: Option<RegistryOptions>,

    ) -> Result<PulledContent, Error> {

        if let Some(layer_type) = layer_type {

            let layer_byte = self

                .pull_layer(name, tag, digest, layer_type, layer_name, options)

                .await?;

            Ok(PulledContent::Layer(layer_byte))

        } else {

            let pack = self.pull_pack(name, tag, digest, options).await?;

            Ok(PulledContent::LyquidPack(pack))

        }

    }

    async fn query(

        &self, name: &str, tag: Option<&str>, digest: Option<&LyquidPackDigest>,

    ) -> Result<LyquidPackMetadata, Error> {

        let reference = self.build_reference(name, tag, digest)?;

        let (_, _, config) = self

            .client

            .pull_manifest_and_config(&reference, &self.auth)

            .await

            .map_err(|e| Error::OciDistributionError(e.to_string()))?;

        let config = config.into_bytes();

        LyquidPackMetadata::from_json(config).map_err(|_| Error::BadMetadata)

    }

}