use atproto_identity::key::{KeyData, KeyType, generate_key, identify_key};
use rusqlite::{OptionalExtension, params};
use serde::{Deserialize, Serialize};
use std::path::Path;
use std::sync::{Arc, Mutex, MutexGuard};

pub const N_CLIENT: &str = "";

#[derive(Serialize, Deserialize, Clone, Debug, PartialEq)]
pub struct Session {
    pub did: String,
    pub client: String,
    pub access_token: String,
    pub refresh_token: String,
    pub dpop_key: String,
    pub pds: String,
    pub issuer: String,
    pub scope: String,
    pub expires_at: i64,
}

#[derive(Clone)]
pub struct Store(pub Arc<Mutex<rusqlite::Connection>>);

const SCHEMA: &str = "
CREATE TABLE IF NOT EXISTS client_keys (client_id TEXT PRIMARY KEY, key TEXT NOT NULL);
CREATE TABLE IF NOT EXISTS requests (state TEXT PRIMARY KEY, request TEXT NOT NULL, created_at INTEGER NOT NULL);
CREATE TABLE IF NOT EXISTS sessions (did TEXT NOT NULL, client TEXT NOT NULL, session TEXT NOT NULL, PRIMARY KEY (did, client));
CREATE TABLE IF NOT EXISTS key_index (did_key TEXT PRIMARY KEY, uri TEXT NOT NULL);
CREATE TABLE IF NOT EXISTS jtis (jti TEXT PRIMARY KEY, exp INTEGER NOT NULL);
";

impl Store {
    pub fn open(path: &Path) -> anyhow::Result<Self> {
        let conn = rusqlite::Connection::open(path)?;
        conn.busy_timeout(std::time::Duration::from_secs(5))?;
        conn.execute_batch(SCHEMA)?;
        Ok(Store(Arc::new(Mutex::new(conn))))
    }

    fn conn(&self) -> MutexGuard<'_, rusqlite::Connection> {
        self.0
            .lock()
            .unwrap_or_else(|poisoned| poisoned.into_inner())
    }

    pub fn client_key(&self, client: &str) -> anyhow::Result<KeyData> {
        let conn = self.conn();
        let existing: Option<String> = conn
            .query_row(
                "SELECT key FROM client_keys WHERE client_id = ?1",
                [client],
                |row| row.get(0),
            )
            .optional()?;
        if let Some(key) = existing {
            return Ok(identify_key(&key)?);
        }
        let key = generate_key(KeyType::P256Private)?;
        conn.execute(
            "INSERT INTO client_keys (client_id, key) VALUES (?1, ?2)",
            params![client, key.to_string()],
        )?;
        Ok(key)
    }

    pub fn put_request(&self, state: &str, request: &str) -> anyhow::Result<()> {
        self.conn().execute(
            "INSERT OR REPLACE INTO requests (state, request, created_at) VALUES (?1, ?2, ?3)",
            params![state, request, chrono::Utc::now().timestamp()],
        )?;
        Ok(())
    }

    pub fn take_request(&self, state: &str) -> anyhow::Result<Option<String>> {
        Ok(self
            .conn()
            .query_row(
                "DELETE FROM requests WHERE state = ?1 RETURNING request",
                [state],
                |row| row.get(0),
            )
            .optional()?)
    }

    pub fn put_session(&self, session: &Session) -> anyhow::Result<()> {
        self.conn().execute(
            "INSERT INTO sessions (did, client, session) VALUES (?1, ?2, ?3)
             ON CONFLICT (did, client) DO UPDATE SET session = excluded.session",
            params![session.did, session.client, serde_json::to_string(session)?],
        )?;
        Ok(())
    }

    pub fn session(&self, did: &str, client: &str) -> anyhow::Result<Option<Session>> {
        let data: Option<String> = self
            .conn()
            .query_row(
                "SELECT session FROM sessions WHERE did = ?1 AND client = ?2",
                [did, client],
                |row| row.get(0),
            )
            .optional()?;
        Ok(data.map(|data| serde_json::from_str(&data)).transpose()?)
    }

    pub fn index_key(&self, did_key: &str, uri: &str) -> anyhow::Result<()> {
        self.conn().execute(
            "INSERT OR REPLACE INTO key_index (did_key, uri) VALUES (?1, ?2)",
            [did_key, uri],
        )?;
        Ok(())
    }

    pub fn key_uri(&self, did_key: &str) -> anyhow::Result<Option<String>> {
        Ok(self
            .conn()
            .query_row(
                "SELECT uri FROM key_index WHERE did_key = ?1",
                [did_key],
                |row| row.get(0),
            )
            .optional()?)
    }

    pub fn remove_key(&self, did_key: &str) -> anyhow::Result<()> {
        self.conn()
            .execute("DELETE FROM key_index WHERE did_key = ?1", [did_key])?;
        Ok(())
    }

    pub fn seen_jti(&self, jti: &str, exp: i64) -> anyhow::Result<bool> {
        let conn = self.conn();
        conn.execute(
            "DELETE FROM jtis WHERE exp < ?1",
            [chrono::Utc::now().timestamp()],
        )?;
        let inserted = conn.execute(
            "INSERT OR IGNORE INTO jtis (jti, exp) VALUES (?1, ?2)",
            params![jti, exp],
        )?;
        Ok(inserted == 0)
    }
}

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

    fn store() -> Store {
        Store::open(Path::new(":memory:")).unwrap()
    }

    fn session(did: &str, client: &str, token: &str) -> Session {
        Session {
            did: did.into(),
            client: client.into(),
            access_token: token.into(),
            refresh_token: "r".into(),
            dpop_key: "k".into(),
            pds: "https://pds.zkwid.test".into(),
            issuer: "https://pds.zkwid.test".into(),
            scope: "atproto".into(),
            expires_at: 1,
        }
    }

    #[test]
    fn client_key_is_stable_per_client() {
        let store = store();
        let a = store.client_key("https://zkwid.test/a").unwrap();
        assert_eq!(a.key_type(), &KeyType::P256Private);
        assert_eq!(
            store
                .client_key("https://zkwid.test/a")
                .unwrap()
                .to_string(),
            a.to_string()
        );
        assert_ne!(
            store
                .client_key("https://zkwid.test/b")
                .unwrap()
                .to_string(),
            a.to_string()
        );
    }

    #[test]
    fn take_request_is_single_use() {
        let store = store();
        store.put_request("s", "{\"a\":1}").unwrap();
        assert_eq!(
            store.take_request("s").unwrap().as_deref(),
            Some("{\"a\":1}")
        );
        assert_eq!(store.take_request("s").unwrap(), None);
        assert_eq!(store.take_request("other").unwrap(), None);
    }

    #[test]
    fn sessions_upsert_by_did_and_client() {
        let store = store();
        assert_eq!(store.session("did:plc:u", N_CLIENT).unwrap(), None);
        store
            .put_session(&session("did:plc:u", N_CLIENT, "a"))
            .unwrap();
        store
            .put_session(&session("did:plc:u", "demo", "b"))
            .unwrap();
        store
            .put_session(&session("did:plc:u", N_CLIENT, "c"))
            .unwrap();
        assert_eq!(
            store.session("did:plc:u", N_CLIENT).unwrap(),
            Some(session("did:plc:u", N_CLIENT, "c"))
        );
        assert_eq!(
            store.session("did:plc:u", "demo").unwrap(),
            Some(session("did:plc:u", "demo", "b"))
        );
        assert_eq!(store.session("did:plc:v", "demo").unwrap(), None);
    }

    #[test]
    fn key_index_round_trip() {
        let store = store();
        assert_eq!(store.key_uri("did:key:z1").unwrap(), None);
        store
            .index_key("did:key:z1", "at://did:plc:u/at.zkwid.key/a")
            .unwrap();
        assert_eq!(
            store.key_uri("did:key:z1").unwrap().as_deref(),
            Some("at://did:plc:u/at.zkwid.key/a")
        );
        store.remove_key("did:key:z1").unwrap();
        assert_eq!(store.key_uri("did:key:z1").unwrap(), None);
    }

    #[test]
    fn seen_jti_detects_replay_and_prunes() {
        let store = store();
        let now = chrono::Utc::now().timestamp();
        assert!(!store.seen_jti("a", now + 60).unwrap());
        assert!(store.seen_jti("a", now + 60).unwrap());
        assert!(!store.seen_jti("old", now - 10).unwrap());
        assert!(!store.seen_jti("old", now + 60).unwrap());
        assert!(store.seen_jti("old", now + 60).unwrap());
    }

    #[test]
    fn persists_across_reopen() {
        let dir = std::env::temp_dir().join(format!("zkwid-n-store-{}", std::process::id()));
        std::fs::create_dir_all(&dir).unwrap();
        let path = dir.join("n.db");
        let key = Store::open(&path)
            .unwrap()
            .client_key("c")
            .unwrap()
            .to_string();
        assert_eq!(
            Store::open(&path)
                .unwrap()
                .client_key("c")
                .unwrap()
                .to_string(),
            key
        );
        std::fs::remove_dir_all(&dir).unwrap();
    }
}