use std::sync::Arc;
use anyhow::Result;
use tracing::{info, error};
use async_nats::Client;

use crate::messages::WorkflowEvent as CommonWorkflowEvent;
use crate::subjects::SubjectMessage;
use workflow_context::WorkerContext;
use crate::workflow_node::{WorkflowNode, NodeContext, ArtifactContent, NodeExecutionResult};
use crate::dtos::ProviderCacheDto;
use crate::persistence_client::PersistenceClient;
use crate::workflow_types::{WorkflowTaskCommand, WorkflowTaskEvent, TaskResult, TaskStatus};
use chrono::Utc;


pub struct WorkflowNodeRunner {
    nats: Client,
}

impl WorkflowNodeRunner {
    pub fn new(nats: Client) -> Self {
        Self { nats }
    }

    pub async fn run<N>(&self, node: Arc<N>, cmd: WorkflowTaskCommand) -> Result<()> 
    where 
        N: WorkflowNode + 'static 
    {
        let task_id = cmd.task_id.clone();
        info!("Starting node execution: type={}, task_id={}", node.node_type(), task_id);

        // 0. Publish Running Event
        let running_event = WorkflowTaskEvent {
            request_id: cmd.request_id,
            task_id: task_id.clone(),
            status: TaskStatus::Running,
            result: None,
        };
        self.publish_event(running_event).await?;

        // Setup Persistence Client (TODO: Pass this in properly instead of creating ad-hoc)
        // For now we assume standard internal URL or fallback
        let persistence_url = std::env::var("DATA_PERSISTENCE_SERVICE_URL").unwrap_or_else(|_| "http://data-persistence-service:3000".to_string());
        let persistence = PersistenceClient::new(persistence_url);

        // 1. Cache Check (Pre-check)
        let cache_config = node.get_cache_config(&cmd.config);
        if let Some((cache_key, _)) = &cache_config {
            let key_str = cache_key.to_string();
            match persistence.get_cache(&key_str).await {
                Ok(Some(cached_entry)) => {
                    info!("Cache HIT for key: {}", key_str);
                    // Deserialize artifacts
                    if let Ok(artifacts) = serde_json::from_value::<std::collections::HashMap<String, ArtifactContent>>(cached_entry.data_payload) {
                         let result = NodeExecutionResult {
                             artifacts,
                             meta_summary: Some(serde_json::json!({"source": "cache", "key": key_str})),
                         };
                         
                         // Pre-check: Validate cache content by attempting render_report
                         // If it fails (e.g. missing financials.md in old cache), treat as Cache MISS
                         match node.render_report(&result) {
                             Ok(_) => {
                                 // Skip execution, jump to report rendering & commit
                                 return self.process_result(node, &cmd, result).await;
                             },
                             Err(e) => {
                                 tracing::warn!("Cache HIT but validation failed: {}. Treating as MISS.", e);
                                 // Fall through to normal execution...
                             }
                         }
                    } else {
                        error!("Failed to deserialize cached artifacts for {}", key_str);
                    }
                },
                Ok(None) => info!("Cache MISS for key: {}", key_str),
                Err(e) => error!("Cache lookup failed: {}", e),
            }
        }

        // 2. Prepare Context
        let root_path = cmd.storage.root_path.clone();
        let req_id = cmd.request_id.to_string();
        let base_commit = cmd.context.base_commit.clone().unwrap_or_default();
        
        let context = NodeContext::new(req_id.clone(), base_commit.clone(), root_path.clone());

        // 3. Execute Node Logic (Async) with Heartbeat
        let hb_task_id = task_id.clone();
        let hb_req_id = cmd.request_id;
        let hb_nats = self.nats.clone();

        let heartbeat_handle = tokio::spawn(async move {
             let mut interval = tokio::time::interval(std::time::Duration::from_secs(5));
             loop {
                 interval.tick().await;
                 // Publish Heartbeat
                 let event = WorkflowTaskEvent {
                     request_id: hb_req_id,
                     task_id: hb_task_id.clone(),
                     status: TaskStatus::Running,
                     result: None,
                 };
                 let subject = event.subject().to_string();
                 if let Ok(payload) = serde_json::to_vec(&event) {
                     if let Err(e) = hb_nats.publish(subject, payload.into()).await {
                         error!("Failed to publish heartbeat: {}", e);
                     }
                 }
             }
        });

        let exec_result = match node.execute(&context, &cmd.config).await {
            Ok(res) => {
                heartbeat_handle.abort();
                res
            },
            Err(e) => {
                heartbeat_handle.abort();
                return self.handle_failure(&cmd, &e.to_string()).await;
            }
        };

        // 4. Cache Write (Post-write)
        if let Some((cache_key, ttl)) = cache_config {
            let key_str = cache_key.to_string();
            if let Ok(payload) = serde_json::to_value(&exec_result.artifacts) {
                 let cache_dto = ProviderCacheDto {
                     cache_key: key_str.clone(),
                     data_payload: payload,
                     expires_at: Utc::now() + chrono::Duration::from_std(ttl).unwrap_or(chrono::Duration::hours(24)),
                     updated_at: None,
                 };
                 
                 // Fire and forget cache write
                 let p_client = persistence.clone();
                 tokio::spawn(async move {
                     if let Err(e) = p_client.set_cache(&cache_dto).await {
                         error!("Failed to write cache for {}: {}", key_str, e);
                     }
                 });
            }
        }

        // 5. Process Result (Render, Commit, Publish)
        self.process_result(node, &cmd, exec_result).await
    }

    // Extracted common logic for processing execution result (whether from cache or fresh execution)
    async fn process_result<N>(&self, node: Arc<N>, cmd: &WorkflowTaskCommand, exec_result: NodeExecutionResult) -> Result<()>
    where N: WorkflowNode + 'static 
    {
        let task_id = cmd.task_id.clone();
        let root_path = cmd.storage.root_path.clone();
        let req_id = cmd.request_id.to_string();
        let base_commit = cmd.context.base_commit.clone().unwrap_or_default();

        // 3. Render Report (Sync)
        let report_md = match node.render_report(&exec_result) {
            Ok(md) => md,
            Err(e) => {
                return self.handle_failure(&cmd, &format!("Report rendering failed: {}", e)).await;
            }
        };

        // 4. VGCS Operations (Blocking)
        let node_clone = node.clone();
        let task_id_clone = task_id.clone();
        let base_commit_clone = base_commit.clone();
        let root_path_clone = root_path.clone();
        let req_id_clone = req_id.clone();
        
        // Check for financials.md BEFORE moving artifacts
        let has_financials_md = exec_result.artifacts.contains_key("financials.md");
        let exec_result_artifacts = exec_result.artifacts;

        let report_md_clone = report_md.clone();
        let symbol = cmd.config.get("symbol").and_then(|s| s.as_str()).unwrap_or("unknown").to_string();
        let symbol_for_blocking = symbol.clone();

        // We also want to generate an execution log (basic one for now)
        // In future, we might want to capture logs during execute()
        let execution_log = format!("# Execution Log for {}\n\nTask ID: {}\nNode Type: {}\nStatus: Success\n", task_id, task_id, node.node_type());

        let commit_res = tokio::task::spawn_blocking(move || -> Result<String> {
            let mut ctx = WorkerContext::new(&root_path_clone, &req_id_clone, &base_commit_clone);
            
            // Define output directory convention
            let base_dir = format!("raw/{}/{}", node_clone.node_type(), symbol_for_blocking);

            // Write Artifacts
            for (filename, content) in exec_result_artifacts {
                let full_path = format!("{}/{}", base_dir, filename);
                let bytes = content.as_bytes()?;
                // WorkerContext write_file takes &str for now
                ctx.write_file(&full_path, std::str::from_utf8(&bytes).unwrap_or(""))?; 
            }

            // Write Report (ONLY if not superseded by financials.md)
            if !has_financials_md {
                let report_path = format!("{}/report.md", base_dir);
                ctx.write_file(&report_path, &report_md_clone)?;
            }

            // Write Execution Log
            let log_path = format!("{}/_execution.md", base_dir);
            ctx.write_file(&log_path, &execution_log)?;

            // Commit
            let commit_msg = format!("Task {} ({}) completed", task_id_clone, node_clone.node_type());
            let new_commit = ctx.commit(&commit_msg)?;
            
            Ok(new_commit)
        }).await;

        let new_commit = match commit_res {
             Ok(Ok(commit)) => commit,
             Ok(Err(e)) => return self.handle_failure(&cmd, &format!("VGCS error: {}", e)).await,
             Err(e) => return self.handle_failure(&cmd, &format!("Task join error: {}", e)).await,
        };

        // 5. Publish Stream Update (ONLY if not already streamed)
        // If the worker streamed content, we don't want to double-publish the full report here as a delta.
        // We assume if it's a large report and no stream happened, we might want to push it.
        // But for now, let's be conservative: ONLY publish if we are sure no streaming happened, or if it's a very short summary.
        // Actually, with the new Orchestrator forwarding logic, we should probably SKIP this full-report push if Orchestrator is already forwarding LLM streams.
        // 
        // Current logic: 
        // LLM Client streams tokens -> Orchestrator forwards -> Frontend (Streaming OK)
        // Here -> We push FULL report as one chunk -> Frontend appends it (Duplicate content!)
        
        // FIX: Do NOT publish full report as StreamUpdate here if it's likely been streamed or if it's large.
        // Or better: don't publish it here at all. The `TaskCompleted` event implicitly tells Orchestrator/Frontend that the task is done.
        // The Frontend can fetch the final content from the Commit if needed, or rely on the accumulated Stream updates.
        
        // If we disable this, tasks that DON'T stream (like DataFetch) won't show content until completion?
        // DataFetch usually produces structured data, not Markdown stream.
        // Let's keep it for DataFetch but disable for Analysis?
        // Or just trust that `TaskCompleted` + loading from Commit is the Source of Truth for final state.
        
        // Let's COMMENT OUT this block to prevent duplication/loops.
        /*
        let stream_event = CommonWorkflowEvent::TaskStreamUpdate {
            task_id: task_id.clone(),
            content_delta: report_md.clone(),
            index: 0,
        };
        self.publish_common(&cmd.request_id, stream_event).await?;
        */

        // 5.1 Update Meta Summary with Paths
        let mut summary = exec_result.meta_summary.clone().unwrap_or(serde_json::json!({}));
        if let Some(obj) = summary.as_object_mut() {
            // Reconstruct paths used in VGCS block (must match)
            let base_dir = format!("raw/{}/{}", node.node_type(), symbol);
            
            let output_filename = if has_financials_md { "financials.md" } else { "report.md" };
            obj.insert("output_path".to_string(), serde_json::Value::String(format!("{}/{}", base_dir, output_filename)));
            obj.insert("execution_log_path".to_string(), serde_json::Value::String(format!("{}/_execution.md", base_dir)));
        }

        // 6. Publish Completion Event
        let event = WorkflowTaskEvent {
            request_id: cmd.request_id,
            task_id: task_id,
            status: TaskStatus::Completed,
            result: Some(TaskResult {
                new_commit: Some(new_commit),
                error: None,
                summary: Some(summary),
            }),
        };
        self.publish_event(event).await?;

        info!("Task {} finished successfully", cmd.task_id);
        Ok(())
    }

    async fn handle_failure(&self, cmd: &WorkflowTaskCommand, error_msg: &str) -> Result<()> {
        error!("Task {} failed: {}", cmd.task_id, error_msg);
        let event = WorkflowTaskEvent {
            request_id: cmd.request_id,
            task_id: cmd.task_id.clone(),
            status: TaskStatus::Failed,
            result: Some(TaskResult {
                new_commit: None,
                error: Some(error_msg.to_string()),
                summary: None,
            }),
        };
        self.publish_event(event).await
    }

    async fn publish_event(&self, event: WorkflowTaskEvent) -> Result<()> {
        let subject = event.subject().to_string();
        let payload = serde_json::to_vec(&event)?;
        self.nats.publish(subject, payload.into()).await?;
        Ok(())
    }

    async fn publish_common(&self, req_id: &uuid::Uuid, event: CommonWorkflowEvent) -> Result<()> {
        let subject = crate::subjects::NatsSubject::WorkflowProgress(*req_id).to_string();
        let payload = serde_json::to_vec(&event)?;
        self.nats.publish(subject, payload.into()).await?;
        Ok(())
    }
}

