use chrono::{DateTime, Utc};
use rusqlite::Connection;
use rusqlite::OptionalExtension;

#[derive(Debug, Clone)]
pub struct Image {
    pub id: i64,
    pub registry: String,
    pub image: String,
    pub expires_at: DateTime<Utc>,
}

#[derive(Debug, Clone)]
pub struct Tag {
    pub id: i64,
    pub image_id: i64,
    pub tag: String,
    pub digest: Option<String>,
}

pub trait Db {
    fn insert_image(&self, image: &mut Image);

    fn get_image(&self, registry: &str, image: &str) -> Option<Image>;
    fn get_expired_images(&self, now: &DateTime<Utc>) -> Vec<Image>;
    fn get_tags_sorted(&self, image_id: i64) -> Vec<Tag>;
    fn insert_tags(&self, tags: &mut [Tag]);
    fn delete_tags(&self, tag_ids: &[i64]);
    fn update_tags(&self, tags: &[Tag]);
    fn set_expires_at(&self, image_id: i64, expires_at: &DateTime<Utc>);
}

pub struct SqliteDb {
    conn: Connection,
}

impl SqliteDb {
    pub fn new(db_path: &std::path::PathBuf) -> Self {
        let conn = Connection::open(db_path).unwrap();
        conn.execute("
            CREATE TABLE IF NOT EXISTS migrations (
                id INTEGER PRIMARY KEY NOT NULL
            )
        ", ()).unwrap();

        let newest_migration: u32 = conn.query_row("
            SELECT MAX(id) FROM migrations
        ", [], |row| row.get::<_, Option<u32>>(0)).unwrap().unwrap_or(0);

        if newest_migration < 1 {
            conn.execute("INSERT INTO migrations(id) VALUES (?1)", (1, )).unwrap();
        }

        if newest_migration < 2 {
            conn.execute("
                CREATE TABLE known_images (
                    id INTEGER PRIMARY KEY NOT NULL,
                    registry TEXT NOT NULL,
                    image TEXT NOT NULL,
                    tag TEXT NOT NULL
                )
            ", ()).unwrap();

            conn.execute("INSERT INTO migrations(id) VALUES (?1)", (2, )).unwrap();
        }

        if newest_migration < 3 {
            conn.execute("
                ALTER TABLE known_images ADD COLUMN
                    discovered DATETIME NOT NULL
            ", ()).unwrap();

            conn.execute("INSERT INTO migrations(id) VALUES (?1)", (3, )).unwrap();
        }

        if newest_migration < 4 {
            conn.execute("INSERT INTO migrations(id) VALUES (?1)", (4, )).unwrap();
        }

        if newest_migration < 5 {
            conn.execute("
                CREATE TABLE images (
                    id INTEGER PRIMARY KEY NOT NULL,
                    registry TEXT NOT NULL,
                    image TEXT NOT NULL,
                    last_checked DATETIME NOT NULL
                )
            ", ()).unwrap();

            conn.execute("
                CREATE UNIQUE INDEX images__registry_image
                    ON images(registry, image)
            ", ()).unwrap();

            conn.execute("INSERT INTO migrations(id) VALUES (?1)", (5, )).unwrap();
        }

        if newest_migration < 6 {
            conn.execute("
                CREATE TABLE tags (
                    id INTEGER PRIMARY KEY,
                    image_id INTEGER NOT NULL REFERENCES images(id),
                    tag TEXT NOT NULL,
                    digest TEXT,
                    fetched_at DATETIME NOT NULL
                )
            ", ()).unwrap();

            conn.execute("
                CREATE UNIQUE INDEX tags__image_id_tag
                    ON tags(image_id, tag)
            ", ()).unwrap();

            conn.execute("INSERT INTO migrations(id) VALUES (?1)", (6, )).unwrap();
        }

        if newest_migration < 7 {
            conn.execute("ALTER TABLE images RENAME COLUMN last_checked TO expires_at", ()).unwrap();
            conn.execute("UPDATE images SET expires_at = datetime(expires_at, '+1440 minutes')", ()).unwrap();
            conn.execute("INSERT INTO migrations(id) VALUES (?1)", (7, )).unwrap();
        }

        if newest_migration < 8 {
            conn.execute("ALTER TABLE tags DROP COLUMN fetched_at", ()).unwrap();
            conn.execute("INSERT INTO migrations(id) VALUES (?1)", (8, )).unwrap();
        }

        return SqliteDb { conn };
    }
}

impl Db for SqliteDb {
    fn get_image(&self, registry: &str, image: &str) -> Option<Image> {
        let _timer = crate::metrics::get().db_query_duration.start_timer();
        return self.conn.query_row("
            SELECT id, registry, image, expires_at FROM images
                WHERE registry = ?1 AND image = ?2
        ", (registry, image), |row| Ok(Image {
            id: row.get(0)?,
            registry: row.get(1)?,
            image: row.get(2)?,
            expires_at: row.get(3)?,
        })).optional().unwrap();
    }

    fn insert_image(&self, image: &mut Image) {
        let _timer = crate::metrics::get().db_query_duration.start_timer();
        self.conn.execute("
            INSERT INTO images(registry, image, expires_at) VALUES (?1, ?2, ?3)
        ", (&image.registry, &image.image, &image.expires_at)).unwrap();
        image.id = self.conn.last_insert_rowid();
    }

    fn get_expired_images(&self, now: &DateTime<Utc>) -> Vec<Image> {
        let _timer = crate::metrics::get().db_query_duration.start_timer();
        let mut stmt = self.conn.prepare("
            SELECT id, registry, image, expires_at FROM images WHERE expires_at <= ?1
        ").unwrap();
        let mut rows = stmt.query((now,)).unwrap();
        let mut images = vec![];
        while let Some(row) = rows.next().unwrap() {
            images.push(Image {
                id: row.get(0).unwrap(),
                registry: row.get(1).unwrap(),
                image: row.get(2).unwrap(),
                expires_at: row.get(3).unwrap(),
            });
        }
        return images;
    }

    fn get_tags_sorted(&self, image_id: i64) -> Vec<Tag> {
        let _timer = crate::metrics::get().db_query_duration.start_timer();
        let mut stmt = self.conn.prepare("SELECT id, image_id, tag, digest FROM tags WHERE image_id = ?1 ORDER BY tag").unwrap();
        let mut rows = stmt.query((image_id,)).unwrap();
        let mut tags = vec![];
        while let Some(row) = rows.next().unwrap() {
            tags.push(Tag {
                id: row.get(0).unwrap(),
                image_id: row.get(1).unwrap(),
                tag: row.get(2).unwrap(),
                digest: row.get(3).unwrap(),
            });
        }
        return tags;
    }

    fn insert_tags(&self, tags: &mut [Tag]) {
        let _timer = crate::metrics::get().db_query_duration.start_timer();
        for tag in tags {
            self.conn.execute("INSERT INTO tags(image_id, tag, digest) VALUES (?1, ?2, ?3)",
                (&tag.image_id, &tag.tag, &tag.digest)).unwrap();
            tag.id = self.conn.last_insert_rowid();
        }
    }

    fn delete_tags(&self, tag_ids: &[i64]) {
        let _timer = crate::metrics::get().db_query_duration.start_timer();
        for id in tag_ids {
            self.conn.execute("DELETE FROM tags WHERE id = ?1", (id,)).unwrap();
        }
    }

    fn update_tags(&self, tags: &[Tag]) {
        let _timer = crate::metrics::get().db_query_duration.start_timer();
        for tag in tags {
            self.conn.execute("UPDATE tags SET tag = ?1, digest = ?2 WHERE id = ?3",
                (&tag.tag, &tag.digest, &tag.id)).unwrap();
        }
    }

    fn set_expires_at(&self, image_id: i64, expires_at: &DateTime<Utc>) {
        let _timer = crate::metrics::get().db_query_duration.start_timer();
        self.conn.execute("UPDATE images SET expires_at = ?1 WHERE id = ?2", (expires_at, image_id)).unwrap();
    }
}

#[cfg(test)]
pub struct StubDb {
    next_id: std::cell::RefCell<i64>,
    images: std::cell::RefCell<Vec<Image>>,
    tags: std::cell::RefCell<Vec<Tag>>,
}

#[cfg(test)]
impl Default for StubDb {
    fn default() -> Self {
        return StubDb {
            next_id: std::cell::RefCell::new(1),
            images: std::cell::RefCell::new(vec![]),
            tags: std::cell::RefCell::new(vec![]),
        };
    }
}

#[cfg(test)]
impl Db for StubDb {
    fn get_image(&self, registry: &str, image: &str) -> Option<Image> {
        for img in self.images.borrow().iter() {
            if img.registry == registry && img.image == image {
                return Some(Image {
                    id: img.id,
                    registry: img.registry.clone(),
                    image: img.image.clone(),
                    expires_at: img.expires_at,
                });
            }
        }
        return None;
    }

    fn insert_image(&self, image: &mut Image) {
        let id = *self.next_id.borrow();
        *self.next_id.borrow_mut() += 1;
        image.id = id;
        self.images.borrow_mut().push(Image {
            id,
            registry: image.registry.clone(),
            image: image.image.clone(),
            expires_at: image.expires_at,
        });
    }

    fn get_expired_images(&self, now: &DateTime<Utc>) -> Vec<Image> {
        let mut result = vec![];
        for img in self.images.borrow().iter() {
            if img.expires_at <= *now {
                result.push(Image {
                    id: img.id,
                    registry: img.registry.clone(),
                    image: img.image.clone(),
                    expires_at: img.expires_at,
                });
            }
        }
        return result;
    }

    fn get_tags_sorted(&self, image_id: i64) -> Vec<Tag> {
        let mut result = vec![];
        for tag in self.tags.borrow().iter() {
            if tag.image_id == image_id {
                result.push(tag.clone());
            }
        }
        result.sort_by(|a, b| a.tag.cmp(&b.tag));
        return result;
    }

    fn insert_tags(&self, tags: &mut [Tag]) {
        for tag in tags {
            let id = *self.next_id.borrow();
            *self.next_id.borrow_mut() += 1;
            tag.id = id;
            self.tags.borrow_mut().push(tag.clone());
        }
    }

    fn delete_tags(&self, tag_ids: &[i64]) {
        self.tags.borrow_mut().retain(|t| !tag_ids.contains(&t.id));
    }

    fn update_tags(&self, tags: &[Tag]) {
        for tag in tags {
            for t in self.tags.borrow_mut().iter_mut() {
                if t.id == tag.id {
                    t.tag = tag.tag.clone();
                    t.digest = tag.digest.clone();
                    break;
                }
            }
        }
    }

    fn set_expires_at(&self, image_id: i64, expires_at: &DateTime<Utc>) {
        for img in self.images.borrow_mut().iter_mut() {
            if img.id == image_id {
                img.expires_at = *expires_at;
                return;
            }
        }
    }
}

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

    fn test_insert_image_returns_incrementing_ids(db: &dyn Db) {
        let now = Utc.with_ymd_and_hms(2000, 1, 1, 0, 0, 0).unwrap();
        let mut img1 = Image { id: 0, registry: "docker.io".into(), image: "nginx".into(), expires_at: now };
        let mut img2 = Image { id: 0, registry: "docker.io".into(), image: "redis".into(), expires_at: now };
        db.insert_image(&mut img1);
        db.insert_image(&mut img2);
        assert!(img2.id > img1.id);
    }

    fn test_get_image_returns_none_for_unknown(db: &dyn Db) {
        assert!(db.get_image("docker.io", "unknown").is_none());
    }

    fn test_get_image_returns_inserted(db: &dyn Db) {
        let now = Utc.with_ymd_and_hms(2000, 1, 1, 0, 0, 0).unwrap();
        let mut img = Image { id: 0, registry: "docker.io".into(), image: "nginx".into(), expires_at: now };
        db.insert_image(&mut img);
        let got = db.get_image("docker.io", "nginx").unwrap();
        assert_eq!(got.id, img.id);
        assert_eq!(got.expires_at, now);
    }

    fn test_get_tags_returns_empty_for_no_tags(db: &dyn Db) {
        let now = Utc.with_ymd_and_hms(2000, 1, 1, 0, 0, 0).unwrap();
        let mut img = Image { id: 0, registry: "docker.io".into(), image: "nginx".into(), expires_at: now };
        db.insert_image(&mut img);
        assert!(db.get_tags_sorted(img.id).is_empty());
    }

    fn test_get_tags_returns_sorted(db: &dyn Db) {
        let now = Utc.with_ymd_and_hms(2000, 1, 1, 0, 0, 0).unwrap();
        let mut img = Image { id: 0, registry: "docker.io".into(), image: "nginx".into(), expires_at: now };
        db.insert_image(&mut img);
        db.insert_tags(&mut [
            Tag { id: 0, image_id: img.id, tag: "2.0".into(), digest: None },
            Tag { id: 0, image_id: img.id, tag: "1.0".into(), digest: None },
            Tag { id: 0, image_id: img.id, tag: "latest".into(), digest: None },
        ]);
        let tags: Vec<String> = db.get_tags_sorted(img.id).into_iter().map(|t| t.tag).collect();
        assert_eq!(tags, vec!["1.0", "2.0", "latest"]);
    }

    fn test_delete_tags_removes_tags(db: &dyn Db) {
        let now = Utc.with_ymd_and_hms(2000, 1, 1, 0, 0, 0).unwrap();
        let mut img = Image { id: 0, registry: "docker.io".into(), image: "nginx".into(), expires_at: now };
        db.insert_image(&mut img);
        let mut tags = [
            Tag { id: 0, image_id: img.id, tag: "1.0".into(), digest: None },
            Tag { id: 0, image_id: img.id, tag: "2.0".into(), digest: None },
        ];
        db.insert_tags(&mut tags);
        db.delete_tags(&[tags[0].id]);
        let result: Vec<String> = db.get_tags_sorted(img.id).into_iter().map(|t| t.tag).collect();
        assert_eq!(result, vec!["2.0"]);
    }

    fn test_set_expires_at(db: &dyn Db) {
        let t1 = Utc.with_ymd_and_hms(2000, 1, 1, 0, 0, 0).unwrap();
        let t2 = Utc.with_ymd_and_hms(2000, 1, 1, 1, 0, 0).unwrap();
        let mut img = Image { id: 0, registry: "docker.io".into(), image: "nginx".into(), expires_at: t1 };
        db.insert_image(&mut img);
        db.set_expires_at(img.id, &t2);
        assert_eq!(db.get_image("docker.io", "nginx").unwrap().expires_at, t2);
    }

    fn test_tag_digest_none_when_unset(db: &dyn Db) {
        let now = Utc.with_ymd_and_hms(2000, 1, 1, 0, 0, 0).unwrap();
        let mut img = Image { id: 0, registry: "docker.io".into(), image: "nginx".into(), expires_at: now };
        db.insert_image(&mut img);
        db.insert_tags(&mut [Tag { id: 0, image_id: img.id, tag: "1.0".into(), digest: None }]);
        assert_eq!(db.get_tags_sorted(img.id)[0].digest, None);
    }

    fn test_update_tags(db: &dyn Db) {
        let now = Utc.with_ymd_and_hms(2000, 1, 1, 0, 0, 0).unwrap();
        let mut img = Image { id: 0, registry: "docker.io".into(), image: "nginx".into(), expires_at: now };
        db.insert_image(&mut img);
        let mut tags = [Tag { id: 0, image_id: img.id, tag: "1.0".into(), digest: None }];
        db.insert_tags(&mut tags);
        tags[0].digest = Some("sha256:abc".into());
        db.update_tags(&tags);
        assert_eq!(db.get_tags_sorted(img.id)[0].digest, Some("sha256:abc".to_string()));
    }

    fn test_get_expired_images(db: &dyn Db) {
        let now = Utc.with_ymd_and_hms(2000, 1, 1, 0, 0, 0).unwrap();
        let past = now - chrono::Duration::days(1);
        let future = now + chrono::Duration::days(1);
        let mut expired = Image { id: 0, registry: "docker.io".into(), image: "nginx".into(), expires_at: past };
        let mut fresh = Image { id: 0, registry: "docker.io".into(), image: "redis".into(), expires_at: future };
        db.insert_image(&mut expired);
        db.insert_image(&mut fresh);
        let result = db.get_expired_images(&now);
        assert_eq!(result.len(), 1);
        assert_eq!(result[0].id, expired.id);
        assert_eq!(result[0].image, "nginx");
    }

    #[test]
    fn conformance_stub_insert_image_ids() {
        test_insert_image_returns_incrementing_ids(&StubDb::default());
    }

    #[test]
    fn conformance_sqlite_insert_image_ids() {
        crate::metrics::init();
        let dir = tempfile::tempdir().unwrap();
        test_insert_image_returns_incrementing_ids(&SqliteDb::new(&dir.path().join("db.sqlite")));
    }

    #[test]
    fn conformance_stub_get_image_unknown() {
        test_get_image_returns_none_for_unknown(&StubDb::default());
    }

    #[test]
    fn conformance_sqlite_get_image_unknown() {
        crate::metrics::init();
        let dir = tempfile::tempdir().unwrap();
        test_get_image_returns_none_for_unknown(&SqliteDb::new(&dir.path().join("db.sqlite")));
    }

    #[test]
    fn conformance_stub_get_image_inserted() {
        test_get_image_returns_inserted(&StubDb::default());
    }

    #[test]
    fn conformance_sqlite_get_image_inserted() {
        crate::metrics::init();
        let dir = tempfile::tempdir().unwrap();
        test_get_image_returns_inserted(&SqliteDb::new(&dir.path().join("db.sqlite")));
    }

    #[test]
    fn conformance_stub_get_tags_empty() {
        test_get_tags_returns_empty_for_no_tags(&StubDb::default());
    }

    #[test]
    fn conformance_sqlite_get_tags_empty() {
        crate::metrics::init();
        let dir = tempfile::tempdir().unwrap();
        test_get_tags_returns_empty_for_no_tags(&SqliteDb::new(&dir.path().join("db.sqlite")));
    }

    #[test]
    fn conformance_stub_get_tags_sorted() {
        test_get_tags_returns_sorted(&StubDb::default());
    }

    #[test]
    fn conformance_sqlite_get_tags_sorted() {
        crate::metrics::init();
        let dir = tempfile::tempdir().unwrap();
        test_get_tags_returns_sorted(&SqliteDb::new(&dir.path().join("db.sqlite")));
    }

    #[test]
    fn conformance_stub_delete_tags() {
        test_delete_tags_removes_tags(&StubDb::default());
    }

    #[test]
    fn conformance_sqlite_delete_tags() {
        crate::metrics::init();
        let dir = tempfile::tempdir().unwrap();
        test_delete_tags_removes_tags(&SqliteDb::new(&dir.path().join("db.sqlite")));
    }

    #[test]
    fn conformance_stub_set_expires_at() {
        test_set_expires_at(&StubDb::default());
    }

    #[test]
    fn conformance_sqlite_set_expires_at() {
        crate::metrics::init();
        let dir = tempfile::tempdir().unwrap();
        test_set_expires_at(&SqliteDb::new(&dir.path().join("db.sqlite")));
    }

    #[test]
    fn conformance_stub_tag_digest_none() {
        test_tag_digest_none_when_unset(&StubDb::default());
    }

    #[test]
    fn conformance_sqlite_tag_digest_none() {
        crate::metrics::init();
        let dir = tempfile::tempdir().unwrap();
        test_tag_digest_none_when_unset(&SqliteDb::new(&dir.path().join("db.sqlite")));
    }

    #[test]
    fn conformance_stub_update_tags() {
        test_update_tags(&StubDb::default());
    }

    #[test]
    fn conformance_sqlite_update_tags() {
        crate::metrics::init();
        let dir = tempfile::tempdir().unwrap();
        test_update_tags(&SqliteDb::new(&dir.path().join("db.sqlite")));
    }

    #[test]
    fn conformance_stub_get_expired_images() {
        test_get_expired_images(&StubDb::default());
    }

    #[test]
    fn conformance_sqlite_get_expired_images() {
        crate::metrics::init();
        let dir = tempfile::tempdir().unwrap();
        test_get_expired_images(&SqliteDb::new(&dir.path().join("db.sqlite")));
    }
}