//! Lineage tracker implementation.

use chrono::Utc;
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::RwLock;
use tracing::{debug, info, instrument};
use uuid::Uuid;

use crate::dag::Task;
use crate::engine::ExecutionStatus;
use crate::{DataValue, Result, state::StateManager};

use super::cache::{LineageCache, LineagePath};
use super::config::LineageConfig;
use super::config::LineageExportFormat;
use super::graph::LineageGraph;
use super::impact::{ImpactAnalysis, ImpactAnalyzer};
use super::node::{LineageEdge, LineageNode};
use super::schema::SchemaEvolutionTracker;
use super::types::{DataSinkType, DataSourceType, LineageEdgeType, LineageNodeType};

/// Data lineage tracker
pub struct LineageTracker {
    /// Configuration
    config: LineageConfig,
    /// State manager for persistence
    state_manager: Arc<StateManager>,
    /// Lineage graph
    lineage_graph: Arc<RwLock<LineageGraph>>,
    /// Schema evolution tracker
    schema_tracker: Arc<RwLock<SchemaEvolutionTracker>>,
    /// Impact analyzer
    impact_analyzer: Arc<ImpactAnalyzer>,
    /// Lineage cache for performance
    lineage_cache: Arc<RwLock<LineageCache>>,
}

impl LineageTracker {
    /// Create a new lineage tracker
    pub fn new(config: LineageConfig, state_manager: Arc<StateManager>) -> Result<Self> {
        let lineage_graph = Arc::new(RwLock::new(LineageGraph::new()));
        let schema_tracker = Arc::new(RwLock::new(SchemaEvolutionTracker::new()));
        let impact_analyzer = Arc::new(ImpactAnalyzer::new(config.impact_analysis.clone()));
        let lineage_cache = Arc::new(RwLock::new(LineageCache::new()));

        Ok(Self {
            config,
            state_manager,
            lineage_graph,
            schema_tracker,
            impact_analyzer,
            lineage_cache,
        })
    }

    /// Track the start of a task
    #[instrument(skip(self, task))]
    pub async fn track_task_start(&self, task_id: &str, task: &Task) -> Result<()> {
        if !self.config.auto_capture {
            return Ok(());
        }

        debug!("Tracking task start: {}", task_id);

        let mut graph = self.lineage_graph.write().await;

        // Create task node
        let task_node = LineageNode {
            node_id: task_id.to_string(),
            node_type: LineageNodeType::Task {
                task_id: task_id.to_string(),
                task_type: format!("{:?}", task.task_type),
            },
            name: task.name.clone(),
            description: task.description.clone(),
            schema: None,
            created_at: Utc::now(),
            modified_at: Utc::now(),
            properties: HashMap::new(),
            tags: vec!["task".to_string()],
        };

        graph.add_node(task_node)?;

        // Track source relationships
        if let Some(ref source) = task.source {
            let source_node_id = format!("source_{}", source.id());
            if !graph.has_node(&source_node_id) {
                let source_node = LineageNode {
                    node_id: source_node_id.clone(),
                    node_type: LineageNodeType::DataSource {
                        source_type: DataSourceType::Other("generic".to_string()),
                        connection_info: HashMap::new(),
                    },
                    name: source.id().to_string(),
                    description: None,
                    schema: None,
                    created_at: Utc::now(),
                    modified_at: Utc::now(),
                    properties: HashMap::new(),
                    tags: vec!["source".to_string()],
                };
                graph.add_node(source_node)?;
            }

            // Add edge from source to task
            graph.add_edge(
                &source_node_id,
                task_id,
                LineageEdge {
                    edge_id: Uuid::new_v4().to_string(),
                    edge_type: LineageEdgeType::DataFlow,
                    strength: 1.0,
                    created_at: Utc::now(),
                    properties: HashMap::new(),
                },
            )?;
        }

        // Track sink relationships
        if let Some(ref sink) = task.sink {
            let sink_node_id = format!("sink_{}", sink.id());
            if !graph.has_node(&sink_node_id) {
                let sink_node = LineageNode {
                    node_id: sink_node_id.clone(),
                    node_type: LineageNodeType::DataSink {
                        sink_type: DataSinkType::Other("generic".to_string()),
                        connection_info: HashMap::new(),
                    },
                    name: sink.id().to_string(),
                    description: None,
                    schema: None,
                    created_at: Utc::now(),
                    modified_at: Utc::now(),
                    properties: HashMap::new(),
                    tags: vec!["sink".to_string()],
                };
                graph.add_node(sink_node)?;
            }

            // Add edge from task to sink
            graph.add_edge(
                task_id,
                &sink_node_id,
                LineageEdge {
                    edge_id: Uuid::new_v4().to_string(),
                    edge_type: LineageEdgeType::DataFlow,
                    strength: 1.0,
                    created_at: Utc::now(),
                    properties: HashMap::new(),
                },
            )?;
        }

        Ok(())
    }

    /// Track the completion of a task
    #[instrument(skip(self))]
    pub async fn track_task_completion(
        &self,
        task_id: &str,
        status: &ExecutionStatus,
        records_processed: u64,
    ) -> Result<()> {
        if !self.config.auto_capture {
            return Ok(());
        }

        debug!(
            "Tracking task completion: {} (status: {:?})",
            task_id, status
        );

        let mut graph = self.lineage_graph.write().await;

        // Update task node with completion information
        if let Some(node) = graph.get_node_mut(task_id) {
            node.modified_at = Utc::now();
            node.properties.insert(
                "execution_status".to_string(),
                DataValue::String(format!("{status:?}")),
            );
            node.properties.insert(
                "records_processed".to_string(),
                DataValue::Int(records_processed as i64),
            );
        }

        Ok(())
    }

    /// Analyze downstream impact of changes to a node
    #[instrument(skip(self))]
    pub async fn analyze_downstream_impact(&self, node_id: &str) -> Result<ImpactAnalysis> {
        info!("Analyzing downstream impact for node: {}", node_id);

        self.impact_analyzer
            .analyze_downstream_impact(node_id, &*self.lineage_graph.read().await)
    }

    /// Analyze upstream dependencies of a node
    #[instrument(skip(self))]
    pub async fn analyze_upstream_dependencies(&self, node_id: &str) -> Result<ImpactAnalysis> {
        info!("Analyzing upstream dependencies for node: {}", node_id);

        self.impact_analyzer
            .analyze_upstream_dependencies(node_id, &*self.lineage_graph.read().await)
    }

    /// Get lineage path between two nodes
    pub async fn get_lineage_path(&self, source: &str, target: &str) -> Result<Vec<LineagePath>> {
        let graph = self.lineage_graph.read().await;
        let mut cache = self.lineage_cache.write().await;

        // Check cache first
        if let Some(paths) = cache.get_path(source, target) {
            return Ok(paths.clone());
        }

        // Calculate paths
        let paths = graph.find_paths(source, target)?;

        // Cache results
        cache.cache_path((source.to_string(), target.to_string()), paths.clone());

        Ok(paths)
    }

    /// Export lineage graph for visualization
    pub async fn export_graph(&self, format: LineageExportFormat) -> Result<String> {
        let graph = self.lineage_graph.read().await;
        graph.export(format)
    }
}