// src/table_script/handlers/dependency_analyzer.rs

use std::collections::HashMap;
use tonic::Status;
use serde_json::{json, Value};

/// Represents the state of a node during dependency graph traversal.
#[derive(Clone, Copy, PartialEq)]
enum NodeState {
    Unvisited,
    Visiting,  // Currently in recursion stack
    Visited,   // Completely processed
}

/// Represents a dependency relationship between tables in Steel scripts.
#[derive(Debug, Clone)]
pub struct Dependency {
    pub target_table: String,
    pub dependency_type: DependencyType,
    pub context: Option<Value>,
}

/// Types of dependencies that can exist between tables in Steel scripts.
#[derive(Debug, Clone)]
pub enum DependencyType {
    /// Direct column access via steel_get_column
    ColumnAccess { column: String },
    /// Indexed column access via steel_get_column_with_index
    IndexedAccess { column: String, index: i64 },
    /// Raw SQL query access via steel_query_sql
    SqlQuery { query_fragment: String },
}

impl DependencyType {
    /// Returns the string representation used in the database.
    pub fn as_str(&self) -> &'static str {
        match self {
            DependencyType::ColumnAccess { .. } => "column_access",
            DependencyType::IndexedAccess { .. } => "indexed_access",
            DependencyType::SqlQuery { .. } => "sql_query",
        }
    }

    /// Generates context JSON for database storage.
    pub fn context_json(&self) -> Value {
        match self {
            DependencyType::ColumnAccess { column } => {
                json!({ "column": column })
            }
            DependencyType::IndexedAccess { column, index } => {
                json!({ "column": column, "index": index })
            }
            DependencyType::SqlQuery { query_fragment } => {
                json!({ "query_fragment": query_fragment })
            }
        }
    }
}

/// Errors that can occur during dependency analysis.
#[derive(Debug)]
pub enum DependencyError {
    CircularDependency {
        cycle_path: Vec<String>,
        involving_script: String
    },
    InvalidTableReference {
        table_name: String,
        script_context: String
    },
    ScriptParseError {
        error: String
    },
    DatabaseError {
        error: String
    },
}

impl From<DependencyError> for Status {
    fn from(error: DependencyError) -> Self {
        match error {
            DependencyError::CircularDependency { cycle_path, involving_script } => {
                Status::failed_precondition(format!(
                    "Circular dependency detected in script for '{}': {}",
                    involving_script,
                    cycle_path.join(" -> ")
                ))
            }
            DependencyError::InvalidTableReference { table_name, script_context } => {
                Status::not_found(format!(
                    "Table '{}' referenced in script '{}' does not exist",
                    table_name, script_context
                ))
            }
            DependencyError::ScriptParseError { error } => {
                Status::invalid_argument(format!("Script parsing failed: {}", error))
            }
            DependencyError::DatabaseError { error } => {
                Status::internal(format!("Database error: {}", error))
            }
        }
    }
}

/// Analyzes Steel scripts to extract table dependencies and validate referential integrity.
///
/// This analyzer identifies how tables reference each other through Steel function calls
/// and ensures that dependency graphs remain acyclic while respecting table link constraints.
pub struct DependencyAnalyzer {
    schema_id: i64,
}

impl DependencyAnalyzer {
    /// Creates a new dependency analyzer for the specified schema.
    pub fn new(schema_id: i64) -> Self {
        Self { schema_id }
    }

    /// Analyzes a Steel script to extract all table dependencies.
    ///
    /// Uses regex patterns to identify function calls that create dependencies:
    /// - `steel_get_column` calls for direct column access
    /// - `steel_get_column_with_index` calls for indexed access
    /// - `steel_query_sql` calls for raw SQL access
    /// - `get-var` calls in transformed scripts
    ///
    /// # Arguments
    /// * `script` - The Steel script code to analyze
    /// * `current_table_name` - Name of the table this script belongs to (for self-references)
    ///
    /// # Returns
    /// * `Ok(Vec<Dependency>)` - List of identified dependencies
    /// * `Err(DependencyError)` - If script parsing fails
    pub fn analyze_script_dependencies(&self, script: &str, current_table_name: &str) -> Result<Vec<Dependency>, DependencyError> {
        let mut dependencies = Vec::new();

        // Extract different types of dependencies using regex patterns
        dependencies.extend(self.extract_function_calls(script)?);
        dependencies.extend(self.extract_sql_dependencies(script)?);
        dependencies.extend(self.extract_get_var_calls(script, current_table_name)?);

        Ok(dependencies)
    }

    /// Extracts Steel function calls that create table dependencies.
    fn extract_function_calls(&self, script: &str) -> Result<Vec<Dependency>, DependencyError> {
        let mut dependencies = Vec::new();

        // Pattern: (steel_get_column "table" "column")
        let column_pattern = regex::Regex::new(r#"\(\s*steel_get_column\s+"([^"]+)"\s+"([^"]+)""#)
            .map_err(|e: regex::Error| DependencyError::ScriptParseError { error: e.to_string() })?;

        for caps in column_pattern.captures_iter(script) {
            let table = caps[1].to_string();
            let column = caps[2].to_string();
            dependencies.push(Dependency {
                target_table: table,
                dependency_type: DependencyType::ColumnAccess { column },
                context: None,
            });
        }

        // Pattern: (steel_get_column_with_index "table" index "column")
        let indexed_pattern = regex::Regex::new(r#"\(\s*steel_get_column_with_index\s+"([^"]+)"\s+(\d+)\s+"([^"]+)""#)
            .map_err(|e: regex::Error| DependencyError::ScriptParseError { error: e.to_string() })?;

        for caps in indexed_pattern.captures_iter(script) {
            let table = caps[1].to_string();
            let index: i64 = caps[2].parse()
                .map_err(|e: std::num::ParseIntError| DependencyError::ScriptParseError { error: e.to_string() })?;
            let column = caps[3].to_string();
            dependencies.push(Dependency {
                target_table: table,
                dependency_type: DependencyType::IndexedAccess { column, index },
                context: None,
            });
        }

        Ok(dependencies)
    }

    /// Extracts get-var calls as dependencies for transformed scripts.
    /// These represent self-references to the current table.
    fn extract_get_var_calls(&self, script: &str, current_table_name: &str) -> Result<Vec<Dependency>, DependencyError> {
        let mut dependencies = Vec::new();

        // Pattern: (get-var "variable")
        let get_var_pattern = regex::Regex::new(r#"\(get-var\s+"([^"]+)"\)"#)
            .map_err(|e: regex::Error| DependencyError::ScriptParseError { error: e.to_string() })?;

        for caps in get_var_pattern.captures_iter(script) {
            let variable_name = caps[1].to_string();
            dependencies.push(Dependency {
                target_table: current_table_name.to_string(),
                dependency_type: DependencyType::ColumnAccess { column: variable_name },
                context: None,
            });
        }

        Ok(dependencies)
    }

    /// Extracts table references from SQL queries in steel_query_sql calls.
    fn extract_sql_dependencies(&self, script: &str) -> Result<Vec<Dependency>, DependencyError> {
        let mut dependencies = Vec::new();

        // Pattern: (steel_query_sql "SELECT ... FROM table ...")
        let sql_pattern = regex::Regex::new(r#"\(\s*steel_query_sql\s+"([^"]+)""#)
            .map_err(|e: regex::Error| DependencyError::ScriptParseError { error: e.to_string() })?;

        for caps in sql_pattern.captures_iter(script) {
            let query = caps[1].to_string();
            let table_refs = self.extract_table_references_from_sql(&query)?;

            for table in table_refs {
                dependencies.push(Dependency {
                    target_table: table.clone(),
                    dependency_type: DependencyType::SqlQuery {
                        query_fragment: query.clone()
                    },
                    context: None,
                });
            }
        }

        Ok(dependencies)
    }

    /// Extracts table names from SQL query text using regex patterns.
    /// Looks for FROM and JOIN clauses to identify table references.
    fn extract_table_references_from_sql(&self, sql: &str) -> Result<Vec<String>, DependencyError> {
        let mut tables = Vec::new();

        // Pattern: FROM table_name or JOIN table_name
        let table_pattern = regex::Regex::new(r#"(?i)\b(?:FROM|JOIN)\s+(?:"([^"]+)"|(\w+))"#)
            .map_err(|e: regex::Error| DependencyError::ScriptParseError { error: e.to_string() })?;

        for caps in table_pattern.captures_iter(sql) {
            let table = caps.get(1)
                .or_else(|| caps.get(2))
                .map(|m| m.as_str().to_string());

            if let Some(table_name) = table {
                tables.push(table_name);
            }
        }

        Ok(tables)
    }

    /// Checks for circular dependencies in the dependency graph.
    ///
    /// This function validates that adding new dependencies won't create cycles
    /// that could lead to infinite loops during script execution. Self-references
    /// are explicitly allowed and filtered out from cycle detection.
    ///
    /// # Arguments
    /// * `tx` - Database transaction for querying existing dependencies
    /// * `table_id` - ID of the table adding new dependencies
    /// * `new_dependencies` - Dependencies to be added
    ///
    /// # Returns
    /// * `Ok(())` - No cycles detected, safe to add dependencies
    /// * `Err(DependencyError)` - Cycle detected or validation failed
    pub async fn check_for_cycles(
        &self,
        tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
        table_id: i64,
        new_dependencies: &[Dependency],
    ) -> Result<(), DependencyError> {
        // Validate that structured table access respects link constraints
        self.validate_link_constraints(tx, table_id, new_dependencies).await?;

        // Build current dependency graph excluding self-references
        let current_deps = sqlx::query!(
            r#"SELECT sd.source_table_id, sd.target_table_id, st.table_name as source_name, tt.table_name as target_name
               FROM script_dependencies sd
               JOIN table_definitions st ON sd.source_table_id = st.id
               JOIN table_definitions tt ON sd.target_table_id = tt.id
               WHERE st.schema_id = $1"#,
            self.schema_id
        )
        .fetch_all(&mut **tx)
        .await
        .map_err(|e| DependencyError::DatabaseError { error: e.to_string() })?;

        let mut graph: HashMap<i64, Vec<i64>> = HashMap::new();
        let mut table_names: HashMap<i64, String> = HashMap::new();

        // Build adjacency list excluding self-references
        for dep in current_deps {
            if dep.source_table_id != dep.target_table_id {
                graph.entry(dep.source_table_id).or_default().push(dep.target_table_id);
            }
            table_names.insert(dep.source_table_id, dep.source_name);
            table_names.insert(dep.target_table_id, dep.target_name);
        }

        // Add new dependencies to test (excluding self-references)
        for dep in new_dependencies {
            let target_id = sqlx::query_scalar!(
                "SELECT id FROM table_definitions WHERE schema_id = $1 AND table_name = $2",
                self.schema_id,
                dep.target_table
            )
            .fetch_optional(&mut **tx)
            .await
            .map_err(|e| DependencyError::DatabaseError { error: e.to_string() })?
            .ok_or_else(|| DependencyError::InvalidTableReference {
                table_name: dep.target_table.clone(),
                script_context: format!("table_id_{}", table_id),
            })?;

            // Only add to cycle detection if not a self-reference
            if table_id != target_id {
                graph.entry(table_id).or_default().push(target_id);
            }

            // Ensure table names are available for error reporting
            if !table_names.contains_key(&table_id) {
                let source_name = sqlx::query_scalar!(
                    "SELECT table_name FROM table_definitions WHERE id = $1",
                    table_id
                )
                .fetch_one(&mut **tx)
                .await
                .map_err(|e| DependencyError::DatabaseError { error: e.to_string() })?;

                table_names.insert(table_id, source_name);
            }
        }

        // Detect cycles using DFS algorithm
        self.detect_cycles_dfs(&graph, &table_names, table_id)?;

        Ok(())
    }

    /// Performs depth-first search to detect cycles in the dependency graph.
    fn dfs_visit(
        &self,
        node: i64,
        states: &mut HashMap<i64, NodeState>,
        graph: &HashMap<i64, Vec<i64>>,
        path: &mut Vec<i64>,
        table_names: &HashMap<i64, String>,
        starting_table: i64,
    ) -> Result<(), DependencyError> {
        states.insert(node, NodeState::Visiting);
        path.push(node);

        if let Some(neighbors) = graph.get(&node) {
            for &neighbor in neighbors {
                if !states.contains_key(&neighbor) {
                    states.insert(neighbor, NodeState::Unvisited);
                }

                match states.get(&neighbor).copied().unwrap_or(NodeState::Unvisited) {
                    NodeState::Visiting => {
                        // Skip self-references as they're allowed
                        if neighbor == node {
                            continue;
                        }

                        // Found a cycle - build the cycle path
                        let cycle_start_idx = path.iter().position(|&x| x == neighbor).unwrap_or(0);
                        let cycle_path: Vec<String> = path[cycle_start_idx..]
                            .iter()
                            .chain(std::iter::once(&neighbor))
                            .map(|&id| table_names.get(&id).cloned().unwrap_or_else(|| id.to_string()))
                            .collect();

                        // Only report as error if cycle involves multiple tables
                        if cycle_path.len() > 2 || (cycle_path.len() == 2 && cycle_path[0] != cycle_path[1]) {
                            let involving_script = table_names.get(&starting_table)
                                .cloned()
                                .unwrap_or_else(|| starting_table.to_string());

                            return Err(DependencyError::CircularDependency {
                                cycle_path,
                                involving_script,
                            });
                        }
                    }
                    NodeState::Unvisited => {
                        self.dfs_visit(neighbor, states, graph, path, table_names, starting_table)?;
                    }
                    NodeState::Visited => {
                        // Already processed, no cycle through this path
                    }
                }
            }
        }

        path.pop();
        states.insert(node, NodeState::Visited);
        Ok(())
    }

    /// Validates that structured table access respects table link constraints.
    ///
    /// # Access Rules
    /// - Self-references are always allowed (table can access its own columns)
    /// - Structured access (steel_get_column functions) requires explicit table links
    /// - Raw SQL access (steel_query_sql) is unrestricted
    ///
    /// # Arguments
    /// * `tx` - Database transaction for querying table links
    /// * `source_table_id` - ID of the table with the script
    /// * `dependencies` - Dependencies to validate
    async fn validate_link_constraints(
        &self,
        tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
        source_table_id: i64,
        dependencies: &[Dependency],
    ) -> Result<(), DependencyError> {
        let current_table_name = sqlx::query_scalar!(
            "SELECT table_name FROM table_definitions WHERE id = $1",
            source_table_id
        )
        .fetch_one(&mut **tx)
        .await
        .map_err(|e| DependencyError::DatabaseError { error: e.to_string() })?;

        // Get all valid linked tables for the source table
        let linked_tables = sqlx::query!(
            r#"SELECT td.table_name, tdl.is_required
               FROM table_definition_links tdl
               JOIN table_definitions td ON tdl.linked_table_id = td.id
               WHERE tdl.source_table_id = $1"#,
            source_table_id
        )
        .fetch_all(&mut **tx)
        .await
        .map_err(|e| DependencyError::DatabaseError { error: e.to_string() })?;

        let mut allowed_tables: std::collections::HashSet<String> = linked_tables
            .into_iter()
            .map(|row| row.table_name)
            .collect();

        // Self-references are always allowed
        allowed_tables.insert(current_table_name.clone());

        // Validate each dependency
        for dep in dependencies {
            match &dep.dependency_type {
                DependencyType::ColumnAccess { column } | DependencyType::IndexedAccess { column, .. } => {
                    // Allow self-references
                    if dep.target_table == current_table_name {
                        continue;
                    }

                    // Check if table is linked
                    if !allowed_tables.contains(&dep.target_table) {
                        return Err(DependencyError::InvalidTableReference {
                            table_name: dep.target_table.clone(),
                            script_context: format!(
                                "Table '{}' is not linked to '{}'. Add a link in the table definition to access '{}' via steel_get_column functions. Column attempted: '{}'. Note: Self-references are always allowed.",
                                dep.target_table,
                                current_table_name,
                                dep.target_table,
                                column
                            ),
                        });
                    }
                }
                DependencyType::SqlQuery { .. } => {
                    // Raw SQL access is unrestricted
                }
            }
        }

        Ok(())
    }

    /// Runs DFS-based cycle detection on the dependency graph.
    fn detect_cycles_dfs(
        &self,
        graph: &HashMap<i64, Vec<i64>>,
        table_names: &HashMap<i64, String>,
        starting_table: i64,
    ) -> Result<(), DependencyError> {
        let mut states: HashMap<i64, NodeState> = HashMap::new();

        // Initialize all nodes as unvisited
        for &node in graph.keys() {
            states.insert(node, NodeState::Unvisited);
        }

        // Run DFS from each unvisited node
        for &node in graph.keys() {
            if states.get(&node) == Some(&NodeState::Unvisited) {
                let mut path = Vec::new();
                if let Err(cycle_error) = self.dfs_visit(
                    node,
                    &mut states,
                    graph,
                    &mut path,
                    table_names,
                    starting_table
                ) {
                    return Err(cycle_error);
                }
            }
        }

        Ok(())
    }

    /// Saves dependencies to the database within an existing transaction.
    ///
    /// This function replaces all existing dependencies for a script with the new set,
    /// ensuring the database reflects the current script analysis results.
    ///
    /// # Arguments
    /// * `tx` - Database transaction for atomic updates
    /// * `script_id` - ID of the script these dependencies belong to
    /// * `table_id` - ID of the table containing the script
    /// * `dependencies` - Dependencies to save
    pub async fn save_dependencies(
        &self,
        tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
        script_id: i64,
        table_id: i64,
        dependencies: &[Dependency],
    ) -> Result<(), DependencyError> {
        // Clear existing dependencies for this script
        sqlx::query!("DELETE FROM script_dependencies WHERE script_id = $1", script_id)
            .execute(&mut **tx)
            .await
            .map_err(|e| DependencyError::DatabaseError { error: e.to_string() })?;

        // Insert new dependencies
        for dep in dependencies {
            let target_id = sqlx::query_scalar!(
                "SELECT id FROM table_definitions WHERE schema_id = $1 AND table_name = $2",
                self.schema_id,
                dep.target_table
            )
            .fetch_optional(&mut **tx)
            .await
            .map_err(|e| DependencyError::DatabaseError { error: e.to_string() })?
            .ok_or_else(|| DependencyError::InvalidTableReference {
                table_name: dep.target_table.clone(),
                script_context: format!("script_id_{}", script_id),
            })?;

            sqlx::query!(
                r#"INSERT INTO script_dependencies
                   (script_id, source_table_id, target_table_id, dependency_type, context_info)
                   VALUES ($1, $2, $3, $4, $5)"#,
                script_id,
                table_id,
                target_id,
                dep.dependency_type.as_str(),
                dep.dependency_type.context_json()
            )
            .execute(&mut **tx)
            .await
            .map_err(|e| DependencyError::DatabaseError { error: e.to_string() })?;
        }

        Ok(())
    }
}
