// Copyright (C) 2026 Doan Luu Duc Tai <doan@luuduct.ai>
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as
// published by the Free Software Foundation, either version 3 of the
// License, or (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.

use crate::store::Store;
use anyhow::{Context, Result, ensure};
use base64::{Engine, prelude::BASE64_STANDARD};
use serde_json::{Value, json};
use std::collections::HashMap;
use std::time::{Duration, Instant};

struct Challenge {
    certificate: String,
    text: String,
    expiry: Instant,
    setup: bool,
}

struct Session {
    identity: String,
    expiry: Instant,
}

pub(crate) struct Service {
    pub(crate) store: Store,
    pub(crate) origin: String,
    setup_token: String,
    challenges: HashMap<String, Challenge>,
    sessions: HashMap<String, Session>,
    window: Instant,
    attempts: u16,
}

pub(crate) fn token() -> Result<String> {
    let mut bytes = [0; 32];
    openssl::rand::rand_bytes(&mut bytes)?;
    Ok(hex::encode(bytes))
}

pub(crate) fn field<'a>(request: &'a Value, key: &str) -> Result<&'a str> {
    request
        .get(key)
        .and_then(Value::as_str)
        .context("Missing string field")
}

fn ciphertext(request: &Value) -> Result<&str> {
    let encoded = field(request, "ciphertext")?;
    ensure!(
        encoded.len() <= 65_536 && BASE64_STANDARD.decode(encoded)?.len() >= 64,
        "Invalid ciphertext length"
    );
    Ok(encoded)
}

fn identifier(value: &str) -> Result<()> {
    ensure!(
        value.len() == 64 && value.bytes().all(|byte| byte.is_ascii_hexdigit()),
        "Invalid identifier"
    );
    Ok(())
}

impl Service {
    pub(crate) fn new(store: Store, origin: String, setup_token: String) -> Self {
        Self {
            store,
            origin,
            setup_token,
            challenges: HashMap::new(),
            sessions: HashMap::new(),
            window: Instant::now(),
            attempts: 0,
        }
    }

    fn owner_id(&self) -> Result<String> {
        stack_chat_crypto::public_identity(&self.store.owner()?.context("Owner not initialized")?)
    }

    fn challenge(&mut self, request: &Value) -> Result<Value> {
        if self.window.elapsed() >= Duration::from_secs(60) {
            self.window = Instant::now();
            self.attempts = 0;
        }
        ensure!(self.attempts < 60, "Authentication rate exceeded");
        self.attempts = self
            .attempts
            .checked_add(1)
            .context("Rate counter overflow")?;
        self.challenges
            .retain(|_, challenge| challenge.expiry > Instant::now());
        ensure!(self.challenges.len() < 128, "Too many pending challenges");
        let certificate = field(request, "public")?;
        let fingerprint = stack_chat_crypto::public_identity(certificate)?;
        let setup = self.store.owner()?.is_none();
        if setup {
            let supplied = field(request, "setup_token")?;
            ensure!(
                self.setup_token.len() == 64
                    && supplied.len() == 64
                    && openssl::memcmp::eq(supplied.as_bytes(), self.setup_token.as_bytes()),
                "Invalid setup credential"
            );
        }
        let nonce = token()?;
        let expires = std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)?
            .as_secs()
            .checked_add(60)
            .context("Clock overflow")?;
        let owner = if setup {
            fingerprint.clone()
        } else {
            self.owner_id()?
        };
        let text = format!(
            "stack-chat-auth/v1\n{}\n{fingerprint}\n{nonce}\n{expires}\n{owner}",
            self.origin
        );
        self.challenges.insert(
            nonce.clone(),
            Challenge {
                certificate: certificate.into(),
                text: text.clone(),
                expiry: Instant::now()
                    .checked_add(Duration::from_secs(60))
                    .context("Clock overflow")?,
                setup,
            },
        );
        Ok(json!({"id":nonce,"challenge":text}))
    }

    fn authenticate(&mut self, request: &Value) -> Result<Value> {
        let challenge = self
            .challenges
            .remove(field(request, "id")?)
            .context("Missing or consumed challenge")?;
        ensure!(challenge.expiry > Instant::now(), "Expired challenge");
        stack_chat_crypto::verify_proof(
            &challenge.certificate,
            &challenge.text,
            field(request, "signature")?,
        )?;
        let identity = stack_chat_crypto::public_identity(&challenge.certificate)?;
        if challenge.setup {
            self.store.install_owner(&challenge.certificate)?;
            self.setup_token.clear();
        }
        let owner = identity == self.owner_id()?;
        if !owner {
            self.store.register(&identity, &challenge.certificate)?;
        }
        self.sessions
            .retain(|_, session| session.expiry > Instant::now());
        ensure!(self.sessions.len() < 256, "Too many sessions");
        let credential = token()?;
        self.sessions.insert(
            credential.clone(),
            Session {
                identity: identity.clone(),
                expiry: Instant::now()
                    .checked_add(Duration::from_mins(30))
                    .context("Clock overflow")?,
            },
        );
        Ok(json!({"token":credential,"owner":owner,"id":identity,"expires_seconds":1800}))
    }

    fn principal(&mut self, request: &Value) -> Result<String> {
        self.sessions
            .retain(|_, session| session.expiry > Instant::now());
        Ok(self
            .sessions
            .get(field(request, "token")?)
            .context("Missing or expired session")?
            .identity
            .clone())
    }

    fn authorized(&self, request: &Value, identity: &str) -> Result<String> {
        let conversation = field(request, "conversation")?;
        identifier(conversation)?;
        ensure!(
            identity == conversation || identity == self.owner_id()?,
            "Forbidden conversation"
        );
        self.store.conversation(conversation)?;
        Ok(conversation.into())
    }

    pub(crate) fn call(&mut self, request: &Value) -> Result<Value> {
        match field(request, "op")? {
            "config" => Ok(
                json!({"owner":self.store.owner()?,"origin":self.origin,"protocol":"stack-chat/v1"}),
            ),
            "challenge" => self.challenge(request),
            "authenticate" => self.authenticate(request),
            _ => self.authenticated(request),
        }
    }

    fn authenticated(&mut self, request: &Value) -> Result<Value> {
        let identity = self.principal(request)?;
        let operation = field(request, "op")?;
        if operation == "lock" {
            self.sessions.remove(field(request, "token")?);
            return Ok(json!({"locked":true}));
        }
        if operation == "inbox" {
            ensure!(identity == self.owner_id()?, "Owner only");
            return self.store.inbox();
        }
        let conversation = self.authorized(request, &identity)?;
        match operation {
            "history" => self.store.messages(
                &conversation,
                request
                    .get("count")
                    .and_then(Value::as_u64)
                    .context("Missing count")?,
                request.get("before").and_then(Value::as_u64).unwrap_or(0),
            ),
            "profile" => {
                ensure!(identity == conversation, "Visitor only");
                self.store.profile(&conversation, ciphertext(request)?)?;
                Ok(json!({"stored":true}))
            }
            "send" => {
                let id = field(request, "id")?;
                identifier(id)?;
                self.store
                    .send(&conversation, id, &identity, ciphertext(request)?)
            }
            _ => anyhow::bail!("Unknown operation"),
        }
    }
}

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

    struct Identity {
        file: String,
        public: String,
        fingerprint: String,
    }

    fn crypto(request: &Value) -> Result<Value> {
        let response =
            stack_chat_crypto::dispatch(&request.to_string()).map_err(anyhow::Error::msg)?;
        Ok(serde_json::from_str(&response)?)
    }

    fn identity(name: &str, owner: &str) -> Result<Identity> {
        let created = crypto(
            &json!({"op":"create","origin":"http://127.0.0.1:9080","name":name,"password":"test-only-strong-password","owner":owner}),
        )?;
        let file = field(&created, "file")?;
        let opened = crypto(
            &json!({"op":"open","origin":"http://127.0.0.1:9080","file":file,"password":"test-only-strong-password"}),
        )?;
        Ok(Identity {
            file: file.into(),
            public: field(&opened, "public")?.into(),
            fingerprint: field(&opened, "fingerprint")?.into(),
        })
    }

    fn authenticate(service: &mut Service, identity: &Identity, setup: &str) -> Result<Value> {
        let challenge = service
            .call(&json!({"op":"challenge","public":identity.public,"setup_token":setup}))?;
        let proof = crypto(
            &json!({"op":"proof","file":identity.file,"password":"test-only-strong-password","origin":service.origin,"challenge":field(&challenge,"challenge")?}),
        )?;
        let request = json!({"op":"authenticate","id":field(&challenge,"id")?,"signature":field(&proof,"signature")?});
        let result = service.call(&request)?;
        ensure!(service.call(&request).is_err(), "Challenge was replayed");
        Ok(result)
    }

    fn reject_expired_and_failed_proofs(service: &mut Service, identity: &Identity) -> Result<()> {
        let challenge = service.call(&json!({"op":"challenge","public":identity.public}))?;
        let challenge_id = field(&challenge, "id")?;
        service
            .challenges
            .get_mut(challenge_id)
            .context("Challenge")?
            .expiry = Instant::now();
        let expired = json!({"op":"authenticate","id":challenge_id,"signature":""});
        ensure!(
            service.call(&expired).is_err(),
            "Expired challenge accepted"
        );
        ensure!(
            !service.challenges.contains_key(challenge_id),
            "Expired challenge retained"
        );
        let invalid = service.call(&json!({"op":"challenge","public":identity.public}))?;
        let invalid_id = field(&invalid, "id")?;
        ensure!(
            service
                .call(&json!({"op":"authenticate","id":invalid_id,"signature":""}))
                .is_err(),
            "Invalid proof accepted"
        );
        ensure!(
            !service.challenges.contains_key(invalid_id),
            "Failed challenge reusable"
        );
        service.attempts = 60;
        ensure!(
            service
                .call(&json!({"op":"challenge","public":identity.public}))
                .is_err(),
            "Challenge rate bound"
        );
        ensure!(
            service.store.install_owner(&identity.public).is_err(),
            "Owner overwritten"
        );
        Ok(())
    }

    #[test]
    fn authorization_replays_history_and_snapshot_recovery() -> Result<()> {
        let directory = TempDir::new()?;
        let path = directory.path().join("chat.sqlite");
        let setup = token()?;
        let mut service = Service::new(
            Store::open(&path)?,
            "http://127.0.0.1:9080".into(),
            setup.clone(),
        );
        let owner = identity("Owner", "")?;
        ensure!(
            service
                .call(&json!({"op":"challenge","public":owner.public,"setup_token":"wrong"}))
                .is_err(),
            "Setup unauthenticated"
        );
        let owner_session = authenticate(&mut service, &owner, &setup)?;
        ensure!(
            owner_session.get("owner") == Some(&json!(true)),
            "Owner role"
        );
        let visitor = identity("Visitor", &owner.public)?;
        let visitor_session = authenticate(&mut service, &visitor, "")?;
        let visitor_token = field(&visitor_session, "token")?;
        let owner_token = field(&owner_session, "token")?;
        let ciphertext = BASE64_STANDARD.encode([42; 128]);
        let message_id = token()?;
        let send = json!({"op":"send","conversation":visitor.fingerprint,"token":visitor_token,"id":message_id,"ciphertext":ciphertext});
        service.call(&send)?;
        service.call(&send)?;
        let conflicting = json!({"op":"send","conversation":visitor.fingerprint,"token":visitor_token,"id":message_id,"ciphertext":BASE64_STANDARD.encode([43;128])});
        ensure!(
            service.call(&conflicting).is_err(),
            "Conflicting ID accepted"
        );
        let empty = service.call(&json!({"op":"history","conversation":visitor.fingerprint,"token":visitor_token,"count":0}))?;
        ensure!(
            empty
                .get("messages")
                .and_then(Value::as_array)
                .is_some_and(Vec::is_empty),
            "N=0"
        );
        ensure!(
            empty
                .get("conversation")
                .and_then(|value| value.get("total"))
                == Some(&json!(1)),
            "Retry duplicated message"
        );
        ensure!(
            service
                .call(&json!({"op":"inbox","token":visitor_token}))
                .is_err(),
            "Visitor read inbox"
        );
        ensure!(service.call(&json!({"op":"history","conversation":owner.fingerprint,"token":visitor_token,"count":1})).is_err(), "Wrong conversation accepted");
        service.call(&json!({"op":"send","conversation":visitor.fingerprint,"token":owner_token,"id":token()?,"ciphertext":ciphertext}))?;
        let latest = service.call(&json!({"op":"history","conversation":visitor.fingerprint,"token":visitor_token,"count":1}))?;
        let sequence = latest
            .get("messages")
            .and_then(Value::as_array)
            .and_then(|messages| messages.first())
            .and_then(|message| message.get("seq"))
            .context("Sequence")?;
        let older = service.call(&json!({"op":"history","conversation":visitor.fingerprint,"token":visitor_token,"count":1,"before":sequence}))?;
        ensure!(
            older
                .get("messages")
                .and_then(Value::as_array)
                .is_some_and(|messages| messages.len() == 1),
            "Pagination"
        );
        ensure!(service.call(&json!({"op":"history","conversation":visitor.fingerprint,"token":visitor_token,"count":51})).is_err(), "History quota");
        reject_expired_and_failed_proofs(&mut service, &visitor)?;
        let snapshot = directory.path().join("snapshot.sqlite");
        service.store.snapshot(&snapshot)?;
        let restored = Store::open(&snapshot)?;
        ensure!(
            restored.owner()? == Some(owner.public),
            "Owner missing after restore"
        );
        ensure!(
            restored.conversation(&visitor.fingerprint)?.get("total") == Some(&json!(2)),
            "Restored count"
        );
        service.call(&json!({"op":"lock","token":visitor_token}))?;
        ensure!(service.call(&json!({"op":"history","conversation":visitor.fingerprint,"token":visitor_token,"count":0})).is_err(), "Locked session accepted");
        let session = service.sessions.get_mut(owner_token).context("Session")?;
        session.expiry = Instant::now();
        ensure!(
            service
                .call(&json!({"op":"inbox","token":owner_token}))
                .is_err(),
            "Expired session accepted"
        );
        Ok(())
    }
}