use std::collections::{HashMap, HashSet};

use std::sync::{Arc, RwLock};

use lyquor_image_store::{Error as ImageRepoError, LyquidImageRepo};

use lyquor_oci::pack::{LyquidPack, LyquidPackDigest};

use lyquor_oci::registry::{OCIRegistryClient, PinnedImage};

use lyquor_primitives::B256;

use thiserror::Error;

#[derive(Debug, Error)]

pub enum Error {

    #[error("Image repository lookup: {0}")]

    ImageRepo(#[from] ImageRepoError),

    #[error("Failed to store resolved image in the local repository.\nDetail: {0}")]

    ImageRepoWrite(String),

    #[error("No image source candidates are known for digest {0:?}.")]

    NoCandidates(B256),

    #[error("Failed to resolve image digest {image_digest:?}.\nAttempts: {attempts}")]

    ResolveFailed { image_digest: B256, attempts: String },

}

pub struct ImageResolver {

    image_repo: Arc<dyn LyquidImageRepo>,

    repo_hints_by_digest: RwLock<HashMap<B256, Vec<String>>>,

    fallback_repository_roots: Vec<String>,

}

impl ImageResolver {

    pub fn new(image_repo: Arc<dyn LyquidImageRepo>, fallback_repository_roots: Vec<String>) -> Self {

        Self {

            image_repo,

            repo_hints_by_digest: RwLock::new(HashMap::new()),

            fallback_repository_roots,

        }

    }

    pub fn image_repo(&self) -> Arc<dyn LyquidImageRepo> {

        self.image_repo.clone()

    }

    pub async fn ensure_local(&self, image_digest: B256) -> Result<(), Error> {

        if self.image_repo.contains(image_digest).await
?0
 {

            return Ok(());

        }

        let candidates = self.repository_candidates(image_digest);

        if candidates.is_empty() {

            return Err(Error::NoCandidates(image_digest));

        }

        let mut failures = Vec::new();

        for candidate in candidates {

            tracing::debug!("Resolving lyquid image digest {image_digest:?} from {candidate}");

            match self.pull_by_digest(&candidate, image_digest).await {

                Ok(pack) => {

                    self.image_repo

                        .put(pack)

                        .await

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

                    self.remember_repository_hint(image_digest, &candidate);

                    return Ok(());

                }

                Err(err) => failures.push(format!("{candidate} ({err})")),

            }

        }

        Err(Error::ResolveFailed {

            image_digest,

            attempts: failures.join(", "),

        })

    }

    pub fn remember_repository_hint(&self, image_digest: B256, repo_hint: &str) {

        if repo_hint.is_empty() {

            return;

        }

        let mut repo_hints_by_digest = self.repo_hints_by_digest.write().expect("repo hint cache poisoned");

        let repo_hints = repo_hints_by_digest.entry(image_digest).or_default();

        if let Some(
pos1
) = repo_hints.iter().position(|candidate| 
candidate1
 == 
repo_hint1
) {

            repo_hints.remove(pos);

        }

        repo_hints.insert(0, repo_hint.to_owned());

    }

    fn repository_candidates(&self, image_digest: B256) -> Vec<String> {

        let cached = self

            .repo_hints_by_digest

            .read()

            .expect("repo hint cache poisoned")

            .get(&image_digest)

            .cloned()

            .unwrap_or_default();

        let mut candidates = Vec::new();

        let mut seen = HashSet::new();

        for candidate in &cached {

            if seen.insert(candidate.clone()) {

                candidates.push(candidate.clone());

            
}0

        }

        for fallback_root in 
&self.fallback_repository_roots3
 {

            if fallback_root.is_empty() {

                continue;

            }

            if seen.insert(fallback_root.clone()) {

                candidates.push(fallback_root.clone());

            
}1

        }

        candidates

    }

    async fn pull_by_digest(&self, repo_hint: &str, image_digest: B256) -> Result<LyquidPack, String> {

        let pinned = PinnedImage::new(repo_hint, LyquidPackDigest::new(image_digest)).map_err(|e| e.to_string())?;

        let oci = OCIRegistryClient::new(pinned.protocol());

        oci.pull_full_pinned(&pinned).await.map_err(|e| e.to_string())

    }

}

#[cfg(test)]

mod tests {

    use super::*;

    use lyquor_image_store::DirectoryRepo;

    use lyquor_test::test;

    use tempfile::TempDir;

    struct TempImageRepo {

        _dir: TempDir,

        repo: Arc<dyn LyquidImageRepo>,

    }

    async fn temp_image_repo() -> TempImageRepo {

        let dir = tempfile::tempdir().expect("temp dir must initialize");

        let repo = Arc::new(

            DirectoryRepo::new(dir.path())

                .await

                .expect("image repo must initialize"),

        );

        TempImageRepo { _dir: dir, repo }

    }

    fn remember_repository_hint(resolver: &ImageResolver, repo_hint: &str, pack: &LyquidPack) {

        resolver.remember_repository_hint(*pack.digest().digest(), repo_hint);

    }

    fn rebuilt_pack(name: &str, description: &str) -> LyquidPack {

        let pack = lyquid_examples::hello();

        let metadata = lyquor_oci::pack::LyquidPackMetadata::new(name, None, Some(description), None, None)

            .with_abi_str(pack.metadata().abi_str.clone());

        LyquidPack::build_with_binary(

            pack.wasm().clone(),

            pack.evm_deployment_bytecode().clone(),

            pack.evm_auxiliary_bytecodes().cloned(),

            pack.assets().cloned(),

            metadata,

        )

    }

    #[test(tokio::test)]

    async fn local_hit_short_circuits_remote_fetch() {

        let pack = rebuilt_pack("hello", "local-hit");

        let image_repo = temp_image_repo().await;

        image_repo.repo.put(pack.clone()).await.expect("put should succeed");

        let resolver = ImageResolver::new(image_repo.repo.clone(), vec!["unused/root".to_owned()]);

        remember_repository_hint(&resolver, "not a valid repo hint", &pack);

        resolver

            .ensure_local(*pack.digest().digest())

            .await

            .expect("local image should resolve");

    }

    #[test(tokio::test)]

    async fn cached_repo_hints_are_tried_before_fallback_roots() {

        let pack = rebuilt_pack("hello", "cached-before-fallback");

        let image_repo = temp_image_repo().await;

        let resolver = ImageResolver::new(image_repo.repo.clone(), vec!["fallback.registry/lyquids".to_owned()]);

        remember_repository_hint(&resolver, "bad.registry/ns/hello", &pack);

        assert_eq!(

            resolver.repository_candidates(*pack.digest().digest()),

            vec![

                "bad.registry/ns/hello".to_owned(),

                "fallback.registry/lyquids".to_owned(),

            ]

        );

    }

    #[test(tokio::test)]

    async fn duplicate_candidates_are_only_attempted_once() {

        let pack = rebuilt_pack("hello", "dedupe");

        let image_repo = temp_image_repo().await;

        let resolver = ImageResolver::new(image_repo.repo.clone(), vec!["same.registry/root/hello".to_owned()]);

        remember_repository_hint(&resolver, "same.registry/root/hello", &pack);

        remember_repository_hint(&resolver, "same.registry/root/hello", &pack);

        assert_eq!(

            resolver.repository_candidates(*pack.digest().digest()),

            vec!["same.registry/root/hello".to_owned()]

        );

    }

    #[test(tokio::test)]

    async fn failed_source_does_not_stop_later_candidates() {

        let pack = rebuilt_pack("hello", "fallback-chain");

        let image_repo = temp_image_repo().await;

        let resolver = ImageResolver::new(

            image_repo.repo.clone(),

            vec!["fallback-one".to_owned(), "fallback-two".to_owned()],

        );

        remember_repository_hint(&resolver, "first.registry/ns/hello", &pack);

        assert_eq!(

            resolver.repository_candidates(*pack.digest().digest()),

            vec![

                "first.registry/ns/hello".to_owned(),

                "fallback-one".to_owned(),

                "fallback-two".to_owned(),

            ]

        );

    }

}