use std::collections::HashMap;
use std::sync::Arc;

use common_contracts::config_models::{
    AnalysisModuleConfig, AnalysisTemplateSets, LlmProvidersConfig,
};
use common_contracts::dtos::{CompanyProfileDto, NewAnalysisResult, TimeSeriesFinancialDto};
use common_contracts::messages::{GenerateReportCommand, ReportGeneratedEvent, ReportFailedEvent};
use common_contracts::subjects::SubjectMessage;
use futures_util::StreamExt;
use petgraph::algo::toposort;
use petgraph::graph::DiGraph;
use tera::{Context, Tera};
use tracing::{info, instrument, error};

use crate::error::{ProviderError, Result};
use crate::llm_client::LlmClient;
use crate::persistence::PersistenceClient;
use crate::state::AppState;
use crate::formatter::format_financials_to_markdown;

#[instrument(skip_all, fields(request_id = %command.request_id, symbol = %command.symbol, template_id = %command.template_id))]
pub async fn run_report_generation_workflow(
    state: Arc<AppState>,
    command: GenerateReportCommand,
) -> Result<()> {
    info!("Starting report generation workflow.");

    // Get channel immediately to ensure we can broadcast errors
    let stream_tx = state.streams.get_or_create_sender(command.request_id);

    let result: Result<()> = async {
        let persistence_client =
            PersistenceClient::new(state.config.data_persistence_service_url.clone());

        // 1. Fetch all necessary data AND configurations in parallel
        info!("Fetching initial data and configurations from persistence service...");
        let (profile, financials, llm_providers, template_sets) =
            fetch_data_and_configs(&persistence_client, command.symbol.as_str(), command.request_id).await
            .map_err(|e| {
                error!("Failed to fetch initial data and configs: {}", e);
                e
            })?;
        info!("Successfully fetched data. Profile: {}, Financials count: {}", profile.name, financials.len());

        // 2. Select the correct template set
        let template_set = template_sets.get(&command.template_id).ok_or_else(|| {
            let err_msg = format!("Analysis template set with ID '{}' not found.", command.template_id);
            error!("{}", err_msg);
            ProviderError::Configuration(err_msg)
        })?;

        // 3. Topologically sort modules to get execution order
        let sorted_modules = sort_modules_by_dependency(&template_set.modules)
            .map_err(|e| {
                error!("Failed to sort analysis modules: {}", e);
                e
            })?;
        
        info!(execution_order = ?sorted_modules, "Successfully determined module execution order.");

        // 4. Execute modules in order
        let mut generated_results: HashMap<String, String> = HashMap::new();

        for module_id in sorted_modules {
            let module_config = template_set.modules.get(&module_id).unwrap();
            
            info!(module_id = %module_id, "All dependencies met. Generating report for module.");

            // Broadcast Module Start
            let _ = stream_tx.send(serde_json::json!({
                "type": "module_start",
                "module_id": module_id
            }).to_string());

            let llm_client = match create_llm_client_for_module(&llm_providers, module_config) {
                Ok(client) => client,
                Err(e) => {
                    error!(module_id = %module_id, "Failed to create LLM client: {}. Skipping module.", e);
                    let err_msg = format!("Error: Failed to create LLM client: {}", e);
                    generated_results.insert(module_id.clone(), err_msg.clone());
                    
                    // Broadcast Error
                    let _ = stream_tx.send(serde_json::json!({
                        "type": "error",
                        "module_id": module_id,
                        "payload": err_msg
                    }).to_string());
                    
                    // Publish Failed Event
                    let fail_event = ReportFailedEvent {
                        request_id: command.request_id,
                        symbol: command.symbol.clone(),
                        module_id: module_id.clone(),
                        error: err_msg.clone(),
                    };
                    if let Ok(payload) = serde_json::to_vec(&fail_event) {
                         if let Err(e) = state.nats.publish(fail_event.subject().to_string(), payload.into()).await {
                             error!("Failed to publish ReportFailedEvent: {}", e);
                         }
                    }
                    
                    continue;
                }
            };

            let mut context = Context::new();
            context.insert("company_name", &profile.name);
            context.insert("ts_code", &command.symbol);

            for dep in &module_config.dependencies {
                if let Some(content) = generated_results.get(dep) {
                    context.insert(dep, content);
                }
            }
            
            // Format financial data into a markdown table
            let formatted_financials = format_financials_to_markdown(&financials);
            context.insert("financial_data", &formatted_financials);

            info!(module_id = %module_id, "Rendering prompt template...");
            let prompt = match Tera::one_off(&module_config.prompt_template, &context, true) {
                Ok(p) => {
                    let p_len = p.len();
                    info!(module_id = %module_id, "Prompt rendered successfully. Length: {} chars", p_len);
                    
                    // Hard Context Limit: 64K chars (~16K tokens)
                    // This is a temporary protection until we have a Deep Research / Summarization module.
                    const MAX_CONTEXT_CHARS: usize = 64_000;
                    if p_len > MAX_CONTEXT_CHARS {
                        let trunc_msg = "\n\n[SYSTEM WARNING: Input data truncated to fit context limits. Analysis may be partial.]";
                        let safe_len = MAX_CONTEXT_CHARS.saturating_sub(trunc_msg.len());
                        let truncated_prompt = format!("{}{}", &p[..safe_len], trunc_msg);
                        
                        tracing::warn!(
                            module_id = %module_id, 
                            "Prompt size ({} chars) exceeded limit ({} chars). Truncated.", 
                            p_len, MAX_CONTEXT_CHARS
                        );
                        truncated_prompt
                    } else {
                        p
                    }
                },
                Err(e) => {
                    let err_msg = format!("Prompt rendering failed: {}", e);
                    error!(module_id = %module_id, "{}", err_msg);
                    generated_results.insert(module_id.clone(), format!("Error: {}", err_msg));
                     // Broadcast Error
                    let _ = stream_tx.send(serde_json::json!({
                        "type": "error",
                        "module_id": module_id,
                        "payload": err_msg
                    }).to_string());

                    // Publish Failed Event
                    let fail_event = ReportFailedEvent {
                        request_id: command.request_id,
                        symbol: command.symbol.clone(),
                        module_id: module_id.clone(),
                        error: err_msg.clone(),
                    };
                    if let Ok(payload) = serde_json::to_vec(&fail_event) {
                         let _ = state.nats.publish(fail_event.subject().to_string(), payload.into()).await;
                    }

                    continue;
                }
            };

            // Streaming Generation
            info!(module_id = %module_id, "Initiating LLM stream...");
            let mut stream = match llm_client.stream_text(prompt).await {
                Ok(s) => s,
                Err(e) => {
                    let err_msg = format!("LLM stream init failed: {}", e);
                    error!(module_id = %module_id, "{}", err_msg);
                    generated_results.insert(module_id.clone(), format!("Error: {}", err_msg));
                     // Broadcast Error
                    let _ = stream_tx.send(serde_json::json!({
                        "type": "error",
                        "module_id": module_id,
                        "payload": err_msg
                    }).to_string());
                    
                    // Publish Failed Event
                    let fail_event = ReportFailedEvent {
                        request_id: command.request_id,
                        symbol: command.symbol.clone(),
                        module_id: module_id.clone(),
                        error: err_msg.clone(),
                    };
                    if let Ok(payload) = serde_json::to_vec(&fail_event) {
                         let _ = state.nats.publish(fail_event.subject().to_string(), payload.into()).await;
                    }

                    continue;
                }
            };
            
            let mut full_content = String::new();
            let mut first_chunk_received = false;

            while let Some(chunk_res) = stream.next().await {
                match chunk_res {
                    Ok(chunk) => {
                        if !chunk.is_empty() {
                            if !first_chunk_received {
                                info!(module_id = %module_id, "Received first chunk from LLM stream.");
                                first_chunk_received = true;
                            }
                            full_content.push_str(&chunk);
                            // Broadcast Content Chunk
                            // Fire and forget - if no subscribers (frontend disconnected), we continue
                            let _ = stream_tx.send(serde_json::json!({
                                "type": "content",
                                "module_id": module_id,
                                "payload": chunk
                            }).to_string());
                        }
                    },
                    Err(e) => {
                        error!(module_id = %module_id, "Stream error: {}", e);
                         let _ = stream_tx.send(serde_json::json!({
                            "type": "error",
                            "module_id": module_id,
                            "payload": format!("Stream error: {}", e)
                        }).to_string());
                        // We might want to break or continue? 
                        // Let's continue to try to finish
                    }
                }
            }

            info!(module_id = %module_id, "Successfully generated content (Length: {}).", full_content.len());

            // Broadcast Module Done
            let _ = stream_tx.send(serde_json::json!({
                "type": "module_done",
                "module_id": module_id
            }).to_string());

            // PERSISTENCE: Only write to DB after full generation
            let result_to_persist = NewAnalysisResult {
                request_id: command.request_id,
                symbol: command.symbol.to_string(),
                template_id: command.template_id.clone(),
                module_id: module_id.clone(),
                content: full_content.clone(),
                meta_data: serde_json::json!({ 
                    "model_id": module_config.model_id,
                    "status": "success"
                }),
            };

            if let Err(e) = persistence_client.create_analysis_result(result_to_persist).await {
                error!(module_id = %module_id, "Failed to persist analysis result: {}", e);
            }

            // Publish Module Completed Event (ReportGenerated)
            // This tells the orchestrator that this module is done.
            let event = ReportGeneratedEvent {
                request_id: command.request_id,
                symbol: command.symbol.clone(),
                module_id: module_id.clone(),
                content_snapshot: Some(full_content.chars().take(200).collect()), // Preview
                model_id: Some(module_config.model_id.clone()),
            };
            
            if let Ok(payload) = serde_json::to_vec(&event) {
                 info!(module_id = %module_id, "Publishing ReportGeneratedEvent");
                 if let Err(e) = state.nats.publish(event.subject().to_string(), payload.into()).await {
                     error!("Failed to publish ReportGeneratedEvent: {}", e);
                 }
            }

            generated_results.insert(module_id.clone(), full_content);
        }
        Ok(())
    }.await;

    if let Err(e) = result {
        error!("Workflow failed: {}", e);
        let _ = stream_tx.send(serde_json::json!({
            "type": "error",
            "module_id": "workflow",
            "payload": format!("Analysis workflow failed: {}", e)
        }).to_string());
    }

    // Broadcast Workflow Done
    let _ = stream_tx.send(serde_json::json!({
        "type": "done"
    }).to_string());

    // Clean up channel with delay to allow clients to connect and receive the final messages
    let streams = state.streams.clone();
    let req_id = command.request_id;
    tokio::spawn(async move {
        tokio::time::sleep(std::time::Duration::from_secs(300)).await; // Keep for 5 minutes to be safe
        streams.remove_channel(&req_id);
    });

    info!("Report generation workflow finished.");
    Ok(())
}

fn sort_modules_by_dependency(
    modules: &HashMap<String, AnalysisModuleConfig>,
) -> Result<Vec<String>> {
    let mut graph = DiGraph::<String, ()>::new();
    let mut node_map = HashMap::new();

    for module_id in modules.keys() {
        let index = graph.add_node(module_id.clone());
        node_map.insert(module_id.clone(), index);
    }

    for (module_id, module_config) in modules {
        if let Some(&module_index) = node_map.get(module_id) {
            for dep in &module_config.dependencies {
                if let Some(&dep_index) = node_map.get(dep) {
                    graph.add_edge(dep_index, module_index, ());
                } else {
                    return Err(ProviderError::Configuration(format!(
                        "Module '{}' has a missing dependency: '{}'",
                        module_id, dep
                    )));
                }
            }
        }
    }

    match toposort(&graph, None) {
        Ok(sorted_nodes) => {
            let sorted_ids = sorted_nodes
                .into_iter()
                .map(|node_index| graph[node_index].clone())
                .collect();
            Ok(sorted_ids)
        }
        Err(cycle) => {
            let cycle_id = graph[cycle.node_id()].clone();
            Err(ProviderError::Configuration(format!(
                "Circular dependency detected in analysis modules. Cycle involves: '{}'",
                cycle_id
            )))
        }
    }
}

fn create_llm_client_for_module(
    llm_providers: &LlmProvidersConfig,
    module_config: &AnalysisModuleConfig,
) -> Result<LlmClient> {
    let provider = llm_providers.get(&module_config.provider_id).ok_or_else(|| {
        ProviderError::Configuration(format!(
            "Provider '{}' not found for module '{}'",
            module_config.provider_id, module_config.name
        ))
    })?;

    let api_url = provider.api_base_url.clone();
    info!("Creating LLM client for module '{}' using provider '{}' with URL: '{}'", module_config.name, module_config.provider_id, api_url);

    Ok(LlmClient::new(
        api_url,
        provider.api_key.clone().into(),
        module_config.model_id.clone(),
        None, // Default timeout
    ))
}

async fn fetch_data_and_configs(
    client: &PersistenceClient,
    symbol: &str,
    request_id: uuid::Uuid,
) -> Result<(
    CompanyProfileDto,
    Vec<TimeSeriesFinancialDto>,
    LlmProvidersConfig,
    AnalysisTemplateSets,
)> {
    info!("fetch_data_and_configs: Starting fetch for request_id: {}", request_id);
    
    // 1. Fetch Configuration (Parallel)
    info!("fetch_data_and_configs: Fetching LLM providers and template sets...");
    let (llm_providers, template_sets) = tokio::try_join!(
        client.get_llm_providers_config(),
        client.get_analysis_template_sets(), 
    ).map_err(|e| {
        error!("fetch_data_and_configs: Failed to fetch configs: {}", e);
        e
    })?;
    info!("fetch_data_and_configs: Configs fetched successfully.");

    // 2. Fetch Session Data (Observational Data Snapshot)
    info!("fetch_data_and_configs: Fetching session data for request_id: {}", request_id);
    let session_data = match client.get_session_data(request_id).await {
        Ok(data) => {
            info!("fetch_data_and_configs: Fetched {} session data items.", data.len());
            data
        },
        Err(e) => {
            error!("Failed to fetch session data for {}: {}", request_id, e);
            vec![] // Treat as empty if failed, though ideally we should fail?
        }
    };

    let mut profile: Option<CompanyProfileDto> = None;
    let mut financials: Vec<TimeSeriesFinancialDto> = Vec::new();

    for item in session_data {
        if item.data_type == "company_profile" {
            if let Ok(p) = serde_json::from_value::<CompanyProfileDto>(item.data_payload.clone()) {
                if profile.is_none() {
                    profile = Some(p);
                }
                // If we have multiple profiles, we currently just take the first one found.
                // Future: We could merge them or provide all to context.
            }
        } else if item.data_type == "financial_statements" {
            if let Ok(mut f_list) = serde_json::from_value::<Vec<TimeSeriesFinancialDto>>(item.data_payload.clone()) {
                // Tag the source for clarity in the report if not already present
                for f in &mut f_list {
                    if f.source.is_none() {
                        f.source = Some(item.provider.clone());
                    }
                }
                financials.extend(f_list);
            }
        }
    }

    // 3. Fallback for Profile (Global Reference)
    if profile.is_none() {
         info!("Profile not found in session data, fetching global profile for {}", symbol);
         if let Ok(p) = client.get_company_profile(symbol).await {
             profile = Some(p);
         }
    }

    // Ensure we have a profile to avoid crashes, or create a dummy one
    let final_profile = profile.unwrap_or_else(|| CompanyProfileDto {
        symbol: symbol.to_string(),
        name: symbol.to_string(),
        industry: None,
        list_date: None,
        additional_info: None,
        updated_at: None,
    });

    Ok((final_profile, financials, llm_providers, template_sets))
}

#[cfg(test)]
mod integration_tests {
    use super::*;
    use crate::config::AppConfig;
    use crate::state::AppState;
    use common_contracts::dtos::SessionDataDto;
    use common_contracts::config_models::{LlmProvider, LlmModel, AnalysisTemplateSet, AnalysisModuleConfig};
    use common_contracts::symbol_utils::{CanonicalSymbol, Market};
    use uuid::Uuid;

    #[tokio::test]
    async fn test_report_generation_flow() {
        if std::env::var("NATS_ADDR").is_err() {
             println!("Skipping integration test (no environment)");
             return;
        }

        // 1. Env & Config
        let config = AppConfig::load().expect("Failed to load config");
        let nats_client = async_nats::connect(config.nats_addr.clone()).await.expect("Failed to connect to NATS in test");
        let state = Arc::new(AppState::new(config.clone(), nats_client));
        let persistence_client = PersistenceClient::new(config.data_persistence_service_url.clone());

        // 2. Setup: LlmProvider
        let api_key = std::env::var("OPENROUTER_API_KEY")
            .expect("OPENROUTER_API_KEY must be set");
        let api_url = std::env::var("OPENROUTER_API_URL")
            .unwrap_or_else(|_| "https://openrouter.ai/api/v1".to_string());

        let mut llm_config = persistence_client.get_llm_providers_config().await.unwrap_or_default();
        llm_config.insert("openrouter".to_string(), LlmProvider {
            api_base_url: api_url,
            api_key: api_key,
            name: "OpenRouter".to_string(),
            models: vec![
                LlmModel {
                    model_id: "google/gemini-flash-1.5".to_string(),
                    name: Some("Gemini Flash".to_string()),
                    is_active: true,
                }
            ],
        });
        persistence_client.update_llm_providers_config(&llm_config).await.expect("Failed to set LLM config");

        // 3. Setup: Template Set
        let mut templates = persistence_client.get_analysis_template_sets().await.unwrap_or_default();
        let mut modules = HashMap::new();
        modules.insert("swot_analysis".to_string(), AnalysisModuleConfig {
            name: "SWOT Analysis".to_string(),
            provider_id: "openrouter".to_string(),
            model_id: "google/gemini-flash-1.5".to_string(),
            prompt_template: "Analyze SWOT for {{ company_name }} ({{ ts_code }}). Keep it brief.".to_string(),
            dependencies: vec![],
        });
        templates.insert("test_template".to_string(), AnalysisTemplateSet {
            name: "Test Template".to_string(),
            modules,
        });
        persistence_client.update_analysis_template_sets(&templates).await.expect("Failed to set Template config");

        // 4. Setup: Session Data (Mocked Financials)
        let request_id = Uuid::new_v4();
        let symbol = CanonicalSymbol::new("AAPL", &Market::US);
        
        // Mock Profile
        let profile_dto = CompanyProfileDto {
            symbol: symbol.to_string(),
            name: "Apple Inc.".to_string(),
            industry: Some("Technology".to_string()),
            list_date: None,
            additional_info: None,
            updated_at: None,
        };
        persistence_client.insert_session_data(&SessionDataDto {
            request_id,
            symbol: symbol.to_string(),
            provider: "mock".to_string(),
            data_type: "company_profile".to_string(),
            data_payload: serde_json::to_value(&profile_dto).unwrap(),
            created_at: None,
        }).await.unwrap();

        // 5. Construct Command
        let cmd = GenerateReportCommand {
            request_id,
            symbol,
            template_id: "test_template".to_string(),
        };

        // 6. Run Workflow
        let result = run_report_generation_workflow(state.clone(), cmd).await;

        // 7. Assert
        assert!(result.is_ok(), "Report generation failed: {:?}", result.err());
    }
}
