//! Server-side HTTP sessions.
//!
//! Mirrors kaya's `SessionMixin` / `SessionStore`:
//!
//! * [`SessionData`] is a dict-like container with modification tracking
//!   (`modified` / `regenerate` / `invalidated` flags) used by the
//!   middleware to decide whether to persist, rotate or delete it;
//! * [`SessionStore`] is the pluggable backend: Redis in production,
//!   in-memory for tests and ephemeral single-process dev;
//! * [`session_middleware`] loads the session from the `session_id` cookie,
//!   exposes it as a [`SessionHandle`] request extension, and persists it
//!   (refreshing the cookie) after the handler ran.
//!
//! Session data is serialized as JSON (the Python server used pickle; old
//! sessions are deliberately not readable, users log in again).

use std::collections::HashMap;
use std::sync::Arc;
use std::time::{Duration, Instant};

use async_trait::async_trait;
use axum::extract::{Request, State};
use axum::middleware::Next;
use axum::response::Response;
use base64::engine::general_purpose::URL_SAFE_NO_PAD;
use base64::Engine;
use rand::RngCore;
use serde_json::{Map, Value};
use tokio::sync::{Mutex, MutexGuard};

use crate::errors::{Error, Result};

/// The session cookie name (kaya's default).
pub const COOKIE_NAME: &str = "session_id";
/// The session idle timeout (kaya's default: 14 days).
pub const MAX_AGE_SECONDS: i64 = 14 * 24 * 60 * 60;
/// Key prefix of the Redis session store. The prefix differs from kaya's
/// (`kaya:session:`) because the payload format changed from pickle to JSON.
pub const REDIS_KEY_PREFIX: &str = "tavolo:session:";

/// Opaque session identifier with 256 bits of entropy, matching Python's
/// `secrets.token_urlsafe(32)` (43 URL-safe base64 characters, no padding).
pub fn new_session_id() -> String {
    let mut bytes = [0u8; 32];
    rand::thread_rng().fill_bytes(&mut bytes);
    URL_SAFE_NO_PAD.encode(bytes)
}

/// Dict-like session container with modification tracking.
#[derive(Debug, Clone, Default)]
pub struct SessionData {
    id: Option<String>,
    old_id: Option<String>,
    data: Map<String, Value>,
    modified: bool,
    regenerate: bool,
    invalidated: bool,
}

impl SessionData {
    pub fn new() -> Self {
        Self::default()
    }

    pub fn with_id(id: impl Into<String>, data: Map<String, Value>) -> Self {
        Self {
            id: Some(id.into()),
            old_id: None,
            data,
            modified: false,
            regenerate: false,
            invalidated: false,
        }
    }

    pub fn id(&self) -> Option<&str> {
        self.id.as_deref()
    }

    pub fn is_modified(&self) -> bool {
        self.modified
    }

    pub fn is_regenerate(&self) -> bool {
        self.regenerate
    }

    pub fn is_invalidated(&self) -> bool {
        self.invalidated
    }

    pub fn get(&self, key: &str) -> Option<&Value> {
        self.data.get(key)
    }

    pub fn insert(&mut self, key: impl Into<String>, value: Value) {
        self.data.insert(key.into(), value);
        self.modified = true;
    }

    pub fn remove(&mut self, key: &str) -> Option<Value> {
        let removed = self.data.remove(key);
        if removed.is_some() {
            self.modified = true;
        }
        removed
    }

    pub fn keys(&self) -> impl Iterator<Item = &String> {
        self.data.keys()
    }

    pub fn data(&self) -> &Map<String, Value> {
        &self.data
    }

    pub fn mark_modified(&mut self) {
        self.modified = true;
    }

    /// Mark the session for ID rotation (session fixation defense); the
    /// middleware mints a new id and deletes the old store entry.
    pub fn regenerate_id(&mut self) {
        self.old_id = self.id.take();
        self.regenerate = true;
        self.modified = true;
    }

    /// Mark the session for deletion and send an expired cookie.
    pub fn invalidate(&mut self) {
        self.invalidated = true;
        self.modified = true;
        self.data.clear();
    }
}

/// Shared, mutable session attached to a request.
#[derive(Clone)]
pub struct SessionHandle(pub Arc<Mutex<SessionData>>);

impl SessionHandle {
    pub fn new(session: SessionData) -> Self {
        Self(Arc::new(Mutex::new(session)))
    }

    pub async fn lock(&self) -> MutexGuard<'_, SessionData> {
        self.0.lock().await
    }
}

/// Pluggable backend for session persistence.
#[async_trait]
pub trait SessionStore: Send + Sync {
    /// Load an existing session, or `None` if unknown or expired. When
    /// `max_age` is given the store slides the server-side expiry.
    async fn load(&self, session_id: &str, max_age: Option<i64>) -> Result<Option<SessionData>>;

    /// Persist the session data under `session_id`.
    async fn save(
        &self,
        session_id: &str,
        data: &Map<String, Value>,
        max_age: Option<i64>,
    ) -> Result<()>;

    /// Remove the session from the store.
    async fn delete(&self, session_id: &str) -> Result<()>;

    /// Number of live sessions (for the `tavolo.active_users` gauge).
    async fn count(&self) -> Result<usize>;
}

// ---------------------------------------------------------------------------
// In-memory store
// ---------------------------------------------------------------------------

struct MemorySession {
    data: Map<String, Value>,
    expires: Option<Instant>,
}

/// Process-local session store used by tests and when Redis is not configured.
#[derive(Default)]
pub struct InMemorySessionStore {
    sessions: Mutex<HashMap<String, MemorySession>>,
}

impl InMemorySessionStore {
    pub fn new() -> Self {
        Self::default()
    }

    fn clock(&self) -> Instant {
        Instant::now()
    }
}

#[async_trait]
impl SessionStore for InMemorySessionStore {
    async fn load(&self, session_id: &str, max_age: Option<i64>) -> Result<Option<SessionData>> {
        let mut sessions = self.sessions.lock().await;
        let Some(entry) = sessions.get_mut(session_id) else {
            return Ok(None);
        };
        let now = self.clock();
        if let Some(expires) = entry.expires {
            if now > expires {
                sessions.remove(session_id);
                return Ok(None);
            }
        }
        if max_age.is_some() {
            entry.expires = Some(now + Duration::from_secs(max_age.unwrap_or(0) as u64));
        }
        Ok(Some(SessionData::with_id(
            session_id.to_string(),
            entry.data.clone(),
        )))
    }

    async fn save(
        &self,
        session_id: &str,
        data: &Map<String, Value>,
        max_age: Option<i64>,
    ) -> Result<()> {
        let mut sessions = self.sessions.lock().await;
        sessions.insert(
            session_id.to_string(),
            MemorySession {
                data: data.clone(),
                expires: max_age.map(|age| self.clock() + Duration::from_secs(age as u64)),
            },
        );
        Ok(())
    }

    async fn delete(&self, session_id: &str) -> Result<()> {
        self.sessions.lock().await.remove(session_id);
        Ok(())
    }

    async fn count(&self) -> Result<usize> {
        let sessions = self.sessions.lock().await;
        let now = self.clock();
        Ok(sessions
            .values()
            .filter(|entry| entry.expires.is_none_or(|expires| expires > now))
            .count())
    }
}

// ---------------------------------------------------------------------------
// Redis store
// ---------------------------------------------------------------------------

/// Redis-backed session store: one JSON value per session under
/// `tavolo:session:<id>`, expiry delegated to the key TTL.
pub struct RedisSessionStore {
    client: redis::Client,
    manager: redis::aio::ConnectionManager,
    prefix: String,
}

impl RedisSessionStore {
    pub async fn new(url: &str) -> Result<Self> {
        let client = redis::Client::open(url)?;
        let manager = redis::aio::ConnectionManager::new(client.clone()).await?;
        Ok(Self {
            client,
            manager,
            prefix: REDIS_KEY_PREFIX.to_string(),
        })
    }

    pub fn with_prefix(mut self, prefix: impl Into<String>) -> Self {
        self.prefix = prefix.into();
        self
    }

    pub fn client(&self) -> redis::Client {
        self.client.clone()
    }

    fn key(&self, session_id: &str) -> String {
        format!("{}{}", self.prefix, session_id)
    }
}

#[async_trait]
impl SessionStore for RedisSessionStore {
    async fn load(&self, session_id: &str, max_age: Option<i64>) -> Result<Option<SessionData>> {
        let mut conn = self.manager.clone();
        let key = self.key(session_id);
        let payload: Option<String> = redis::cmd("GET").arg(&key).query_async(&mut conn).await?;
        let Some(payload) = payload else {
            return Ok(None);
        };
        if let Some(max_age) = max_age {
            redis::cmd("EXPIRE")
                .arg(&key)
                .arg(max_age)
                .query_async::<i64>(&mut conn)
                .await?;
        }
        let data: Map<String, Value> = serde_json::from_str(&payload)?;
        Ok(Some(SessionData::with_id(session_id.to_string(), data)))
    }

    async fn save(
        &self,
        session_id: &str,
        data: &Map<String, Value>,
        max_age: Option<i64>,
    ) -> Result<()> {
        let mut conn = self.manager.clone();
        let payload = serde_json::to_string(data)?;
        let mut cmd = redis::cmd("SET");
        cmd.arg(self.key(session_id)).arg(payload);
        if let Some(max_age) = max_age {
            cmd.arg("EX").arg(max_age);
        }
        cmd.query_async::<()>(&mut conn).await?;
        Ok(())
    }

    async fn delete(&self, session_id: &str) -> Result<()> {
        let mut conn = self.manager.clone();
        redis::cmd("DEL")
            .arg(self.key(session_id))
            .query_async::<i64>(&mut conn)
            .await?;
        Ok(())
    }

    async fn count(&self) -> Result<usize> {
        let mut conn = self.manager.clone();
        let mut cursor: u64 = 0;
        let mut total = 0usize;
        loop {
            let (next, keys): (u64, Vec<String>) = redis::cmd("SCAN")
                .arg(cursor)
                .arg("MATCH")
                .arg(format!("{}*", self.prefix))
                .arg("COUNT")
                .arg(500)
                .query_async(&mut conn)
                .await?;
            total += keys.len();
            cursor = next;
            if cursor == 0 {
                break;
            }
        }
        Ok(total)
    }
}

/// Convenience alias used by the application.
pub type SharedSessionStore = Arc<dyn SessionStore>;

// ---------------------------------------------------------------------------
// Cookie helpers
// ---------------------------------------------------------------------------

/// Return the value of `name` from a `Cookie` header, if present.
pub fn parse_cookie_value(header_value: &str, name: &str) -> Option<String> {
    for part in header_value.split(';') {
        let part = part.trim();
        if let Some((key, value)) = part.split_once('=') {
            if key.trim() == name {
                return Some(value.trim().trim_matches('"').to_string());
            }
        }
    }
    None
}

/// Build a `Set-Cookie` value the way Python's `SimpleCookie` renders one,
/// i.e. sorted attributes: `HttpOnly; Max-Age=n; Path=/; SameSite=Lax`.
pub fn format_set_cookie(
    name: &str,
    value: &str,
    path: &str,
    max_age: Option<i64>,
    httponly: bool,
    secure: bool,
    samesite: Option<&str>,
) -> String {
    let mut out = format!("{name}={value}");
    if httponly {
        out.push_str("; HttpOnly");
    }
    if let Some(max_age) = max_age {
        out.push_str(&format!("; Max-Age={max_age}"));
    }
    out.push_str(&format!("; Path={path}"));
    if let Some(samesite) = samesite {
        out.push_str(&format!("; SameSite={samesite}"));
    }
    if secure {
        out.push_str("; Secure");
    }
    out
}

fn extract_session_id(headers: &axum::http::HeaderMap) -> Option<String> {
    let mut found = None;
    for value in headers.get_all(axum::http::header::COOKIE) {
        if let Ok(value) = value.to_str() {
            if let Some(id) = parse_cookie_value(value, COOKIE_NAME) {
                found = Some(id);
                break;
            }
        }
    }
    found
}

/// Load the session, expose it as a [`SessionHandle`] extension and persist
/// it (refreshing the cookie) after the handler ran.
///
/// Mirrors kaya's `SessionMixin`: an existing session is always re-saved and
/// its cookie re-sent, so the client and server expiry windows slide on
/// every request.
pub async fn session_middleware(
    State(store): State<SharedSessionStore>,
    mut request: Request,
    next: Next,
) -> Response {
    // Static assets never reach the application in the Python deployment
    // (Granian serves them): skip the session round-trip for them.
    let path = request.uri().path();
    if path.starts_with("/static") || path.starts_with("/assets") {
        return next.run(request).await;
    }
    let session_id = extract_session_id(request.headers());
    let session = match session_id.as_deref() {
        Some(id) => match store.load(id, Some(MAX_AGE_SECONDS)).await {
            Ok(Some(session)) => session,
            Ok(None) => SessionData::new(),
            Err(err) => {
                tracing::error!("session load failed: {err}");
                SessionData::new()
            }
        },
        None => SessionData::new(),
    };
    let handle = SessionHandle::new(session);
    request.extensions_mut().insert(handle.clone());

    let mut response = next.run(request).await;

    let mut session = handle.lock().await;
    let mut cookie: Option<String> = None;
    if session.invalidated {
        let id = session.old_id.clone().or_else(|| session.id.clone());
        if let Some(id) = &id {
            if let Err(err) = store.delete(id).await {
                tracing::error!("session delete failed: {err}");
            }
        }
        cookie = Some(format_set_cookie(
            COOKIE_NAME,
            id.as_deref().unwrap_or(""),
            "/",
            Some(0),
            true,
            false,
            Some("Lax"),
        ));
    } else if session.regenerate {
        if let Some(old_id) = session.old_id.clone() {
            if let Err(err) = store.delete(&old_id).await {
                tracing::error!("session delete failed: {err}");
            }
            session.old_id = None;
        }
        if session.id.is_none() {
            session.id = Some(new_session_id());
        }
        session.regenerate = false;
        if let Some(id) = session.id.clone() {
            if let Err(err) = store.save(&id, session.data(), Some(MAX_AGE_SECONDS)).await {
                tracing::error!("session save failed: {err}");
            }
            cookie = Some(format_set_cookie(
                COOKIE_NAME,
                &id,
                "/",
                Some(MAX_AGE_SECONDS),
                true,
                false,
                Some("Lax"),
            ));
        }
    } else {
        if session.id.is_none() && session.is_modified() {
            session.id = Some(new_session_id());
        }
        if let Some(id) = session.id.clone() {
            if let Err(err) = store.save(&id, session.data(), Some(MAX_AGE_SECONDS)).await {
                tracing::error!("session save failed: {err}");
            }
            cookie = Some(format_set_cookie(
                COOKIE_NAME,
                &id,
                "/",
                Some(MAX_AGE_SECONDS),
                true,
                false,
                Some("Lax"),
            ));
        }
    }
    drop(session);

    if let Some(cookie) = cookie {
        if let Ok(value) = axum::http::HeaderValue::from_str(&cookie) {
            response
                .headers_mut()
                .append(axum::http::header::SET_COOKIE, value);
        }
    }
    response
}

/// Errors from the session layer surfaced as platform errors.
pub fn session_error(message: impl Into<String>) -> Error {
    Error::Internal(message.into())
}