//! Node enrollment and dynamic registration system

use chrono::{DateTime, Utc};
use rusqlite::{params, Connection, Result as SqliteResult};
use serde::{Deserialize, Serialize};
use std::sync::{Arc, Mutex};
use thiserror::Error;
use uuid::Uuid;

#[derive(Error, Debug)]
pub enum EnrollmentError {
    #[error("database error: {0}")]
    Database(#[from] rusqlite::Error),
    #[error("node not found")]
    NotFound,
    #[error("invalid state transition")]
    InvalidTransition,
    #[error("crypto error: {0}")]
    Crypto(String),
}

/// Node enrollment state machine
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum EnrollmentState {
    Pending,
    Approved,
    Active,
    Rejected,
    Suspended,
}

impl EnrollmentState {
    pub fn as_str(&self) -> &'static str {
        match self {
            Self::Pending => "pending",
            Self::Approved => "approved",
            Self::Active => "active",
            Self::Rejected => "rejected",
            Self::Suspended => "suspended",
        }
    }

    pub fn from_str(s: &str) -> Option<Self> {
        match s {
            "pending" => Some(Self::Pending),
            "approved" => Some(Self::Approved),
            "active" => Some(Self::Active),
            "rejected" => Some(Self::Rejected),
            "suspended" => Some(Self::Suspended),
            _ => None,
        }
    }
}

/// Node enrollment request
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EnrollmentRequest {
    pub token: String,           // OTP enrollment token (required)
    pub node_id: Uuid,
    pub name: String,
    pub address: String,         // Public IP:port
    pub country_code: String,
    pub hardware_id: String,     // Unique hardware fingerprint
    pub region: String,          // Geographic region
    pub requested_at: DateTime<Utc>,
}

/// Enrolled node record
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EnrolledNode {
    pub node_id: Uuid,
    pub name: String,
    pub address: String,
    pub country_code: String,
    pub hardware_id: String,
    pub region: String,
    pub state: EnrollmentState,
    pub node_psk: Option<String>,  // Generated PSK for approved nodes
    pub requested_at: DateTime<Utc>,
    pub approved_at: Option<DateTime<Utc>>,
    pub activated_at: Option<DateTime<Utc>>,
}

impl EnrolledNode {
    /// Check if node can transition to new state
    pub fn can_transition_to(&self, new_state: EnrollmentState) -> bool {
        match (self.state, new_state) {
            (EnrollmentState::Pending, EnrollmentState::Approved) => true,
            (EnrollmentState::Pending, EnrollmentState::Rejected) => true,
            (EnrollmentState::Approved, EnrollmentState::Active) => true,
            (EnrollmentState::Approved, EnrollmentState::Rejected) => true,
            (EnrollmentState::Active, EnrollmentState::Suspended) => true,
            (EnrollmentState::Suspended, EnrollmentState::Active) => true,
            _ => false,
        }
    }
}

/// Node enrollment registry with SQLite backend
pub struct EnrollmentRegistry {
    conn: Arc<Mutex<Connection>>,
}

impl EnrollmentRegistry {
    /// Create new enrollment registry with SQLite database
    pub fn new(db_path: impl AsRef<std::path::Path>) -> Result<Self, EnrollmentError> {
        let conn = Connection::open(db_path)?;
        
        conn.execute(
            "CREATE TABLE IF NOT EXISTS enrollments (
                node_id TEXT PRIMARY KEY,
                name TEXT NOT NULL,
                address TEXT NOT NULL,
                country_code TEXT NOT NULL,
                hardware_id TEXT NOT NULL,
                region TEXT NOT NULL,
                state TEXT NOT NULL,
                node_psk TEXT,
                requested_at TEXT NOT NULL,
                approved_at TEXT,
                activated_at TEXT
            )",
            [],
        )?;

        Ok(Self {
            conn: Arc::new(Mutex::new(conn)),
        })
    }

    /// Submit enrollment request (server wants to join network)
    pub fn submit_request(&self, req: EnrollmentRequest) -> Result<(), EnrollmentError> {
        let conn = self.conn.lock().unwrap();
        
        conn.execute(
            "INSERT INTO enrollments (
                node_id, name, address, country_code, hardware_id, region,
                state, node_psk, requested_at, approved_at, activated_at
            ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, NULL, ?8, NULL, NULL)",
            params![
                req.node_id.to_string(),
                req.name,
                req.address,
                req.country_code,
                req.hardware_id,
                req.region,
                EnrollmentState::Pending.as_str(),
                req.requested_at.to_rfc3339(),
            ],
        )?;

        tracing::info!("Enrollment request submitted for node {}", req.node_id);
        Ok(())
    }

    /// Get enrollment by node ID
    pub fn get(&self, node_id: &Uuid) -> Result<Option<EnrolledNode>, EnrollmentError> {
        let conn = self.conn.lock().unwrap();
        
        let mut stmt = conn.prepare(
            "SELECT node_id, name, address, country_code, hardware_id, region,
                    state, node_psk, requested_at, approved_at, activated_at
             FROM enrollments WHERE node_id = ?1"
        )?;

        let mut rows = stmt.query(params![node_id.to_string()])?;
        
        if let Some(row) = rows.next()? {
            Ok(Some(Self::row_to_node(row)?))
        } else {
            Ok(None)
        }
    }

    /// List all enrollments by state
    pub fn list_by_state(&self, state: EnrollmentState) -> Result<Vec<EnrolledNode>, EnrollmentError> {
        let conn = self.conn.lock().unwrap();
        
        let mut stmt = conn.prepare(
            "SELECT node_id, name, address, country_code, hardware_id, region,
                    state, node_psk, requested_at, approved_at, activated_at
             FROM enrollments WHERE state = ?1"
        )?;

        let rows = stmt.query_map(params![state.as_str()], |row| {
            Ok(Self::row_to_node(row).unwrap())
        })?;

        Ok(rows.collect::<SqliteResult<Vec<_>>>()?)
    }

    /// Approve enrollment request and generate Node PSK with IP allocation
    pub fn approve(&self, node_id: &Uuid, assigned_ip: &str) -> Result<String, EnrollmentError> {
        let conn = self.conn.lock().unwrap();
        
        // Generate unique PSK for this node
        let node_psk = Self::generate_node_psk();
        let now = Utc::now().to_rfc3339();

        // First check if we need to add assigned_ip column
        let has_ip_column = conn.query_row(
            "SELECT COUNT(*) FROM pragma_table_info('enrollments') WHERE name='assigned_ip'",
            [],
            |row| row.get::<_, i32>(0)
        ).unwrap_or(0) > 0;

        if !has_ip_column {
            conn.execute(
                "ALTER TABLE enrollments ADD COLUMN assigned_ip TEXT",
                [],
            )?;
        }

        let updated = conn.execute(
            "UPDATE enrollments 
             SET state = ?1, node_psk = ?2, approved_at = ?3, assigned_ip = ?4
             WHERE node_id = ?5 AND state = ?6",
            params![
                EnrollmentState::Approved.as_str(),
                node_psk,
                now,
                assigned_ip,
                node_id.to_string(),
                EnrollmentState::Pending.as_str(),
            ],
        )?;

        if updated == 0 {
            return Err(EnrollmentError::InvalidTransition);
        }

        tracing::info!("Node {} approved with IP {}", node_id, assigned_ip);
        Ok(node_psk)
    }

    /// Reject enrollment request
    pub fn reject(&self, node_id: &Uuid) -> Result<(), EnrollmentError> {
        let conn = self.conn.lock().unwrap();
        
        let updated = conn.execute(
            "UPDATE enrollments SET state = ?1 WHERE node_id = ?2 AND state = ?3",
            params![
                EnrollmentState::Rejected.as_str(),
                node_id.to_string(),
                EnrollmentState::Pending.as_str(),
            ],
        )?;

        if updated == 0 {
            return Err(EnrollmentError::InvalidTransition);
        }

        tracing::info!("Node {} rejected", node_id);
        Ok(())
    }

    /// Activate approved node (node has successfully connected)
    pub fn activate(&self, node_id: &Uuid) -> Result<(), EnrollmentError> {
        let conn = self.conn.lock().unwrap();
        
        let now = Utc::now().to_rfc3339();
        
        let updated = conn.execute(
            "UPDATE enrollments 
             SET state = ?1, activated_at = ?2
             WHERE node_id = ?3 AND state = ?4",
            params![
                EnrollmentState::Active.as_str(),
                now,
                node_id.to_string(),
                EnrollmentState::Approved.as_str(),
            ],
        )?;

        if updated == 0 {
            return Err(EnrollmentError::InvalidTransition);
        }

        tracing::info!("Node {} activated", node_id);
        Ok(())
    }

    /// Suspend active node
    pub fn suspend(&self, node_id: &Uuid) -> Result<(), EnrollmentError> {
        let conn = self.conn.lock().unwrap();
        
        let updated = conn.execute(
            "UPDATE enrollments SET state = ?1 WHERE node_id = ?2 AND state = ?3",
            params![
                EnrollmentState::Suspended.as_str(),
                node_id.to_string(),
                EnrollmentState::Active.as_str(),
            ],
        )?;

        if updated == 0 {
            return Err(EnrollmentError::InvalidTransition);
        }

        tracing::info!("Node {} suspended", node_id);
        Ok(())
    }

    /// Generate secure Node PSK (32 bytes hex)
    fn generate_node_psk() -> String {
        use rand::RngCore;
        let mut psk = [0u8; 32];
        rand::thread_rng().fill_bytes(&mut psk);
        hex::encode(psk)
    }

    /// Convert SQLite row to EnrolledNode
    fn row_to_node(row: &rusqlite::Row) -> SqliteResult<EnrolledNode> {
        Ok(EnrolledNode {
            node_id: Uuid::parse_str(&row.get::<_, String>(0)?).unwrap(),
            name: row.get(1)?,
            address: row.get(2)?,
            country_code: row.get(3)?,
            hardware_id: row.get(4)?,
            region: row.get(5)?,
            state: EnrollmentState::from_str(&row.get::<_, String>(6)?).unwrap(),
            node_psk: row.get(7)?,
            requested_at: DateTime::parse_from_rfc3339(&row.get::<_, String>(8)?)
                .unwrap()
                .into(),
            approved_at: row
                .get::<_, Option<String>>(9)?
                .and_then(|s| DateTime::parse_from_rfc3339(&s).ok())
                .map(|dt| dt.into()),
            activated_at: row
                .get::<_, Option<String>>(10)?
                .and_then(|s| DateTime::parse_from_rfc3339(&s).ok())
                .map(|dt| dt.into()),
        })
    }
}

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

    #[test]
    fn test_state_transitions() {
        let node = EnrolledNode {
            node_id: Uuid::new_v4(),
            name: "test".into(),
            address: "1.2.3.4:8443".into(),
            country_code: "US".into(),
            hardware_id: "hwid".into(),
            region: "us-west".into(),
            state: EnrollmentState::Pending,
            node_psk: None,
            requested_at: Utc::now(),
            approved_at: None,
            activated_at: None,
        };

        assert!(node.can_transition_to(EnrollmentState::Approved));
        assert!(node.can_transition_to(EnrollmentState::Rejected));
        assert!(!node.can_transition_to(EnrollmentState::Active));
    }

    #[test]
    fn test_enrollment_workflow() -> Result<(), EnrollmentError> {
        let registry = EnrollmentRegistry::new(":memory:")?;
        let node_id = Uuid::new_v4();

        // Submit request
        let req = EnrollmentRequest {
            token: "TEST_TOKEN_123".into(),
            node_id,
            name: "test-node".into(),
            address: "1.2.3.4:8443".into(),
            country_code: "US".into(),
            hardware_id: "test-hw".into(),
            region: "us-west".into(),
            requested_at: Utc::now(),
        };
        registry.submit_request(req)?;

        // Check pending
        let pending = registry.list_by_state(EnrollmentState::Pending)?;
        assert_eq!(pending.len(), 1);

        // Approve
        let psk = registry.approve(&node_id, "10.42.0.2")?;
        assert_eq!(psk.len(), 64); // 32 bytes hex

        // Check approved
        let approved = registry.list_by_state(EnrollmentState::Approved)?;
        assert_eq!(approved.len(), 1);

        // Activate
        registry.activate(&node_id)?;

        // Check active
        let active = registry.list_by_state(EnrollmentState::Active)?;
        assert_eq!(active.len(), 1);

        Ok(())
    }
}
