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; // Assuming this command is defined
use petgraph::algo::toposort;
use petgraph::graph::DiGraph;
use tera::{Context, Tera};
use tracing::{info, warn, instrument, error};
use uuid::Uuid;

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

#[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.");

    let persistence_client =
        PersistenceClient::new(state.config.data_persistence_service_url.clone());

    // 1. Fetch all necessary data AND configurations in parallel
    let (profile, financials, llm_providers, template_sets) =
        match fetch_data_and_configs(&persistence_client, &command.symbol).await {
            Ok(data) => data,
            Err(e) => {
                error!("Failed to fetch initial data and configs: {}", e);
                return Err(e);
            }
        };

    // 2. Select the correct template set
    let template_set = match template_sets.get(&command.template_id) {
        Some(ts) => ts,
        None => {
            let err_msg = format!("Analysis template set with ID '{}' not found.", command.template_id);
            error!("{}", err_msg);
            return Err(ProviderError::Configuration(err_msg));
        }
    };

    // 3. Topologically sort modules to get execution order
    let sorted_modules = match sort_modules_by_dependency(&template_set.modules) {
        Ok(order) => order,
        Err(e) => {
            error!("Failed to sort analysis modules: {}", e);
            return Err(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(); // Should not fail due to sorting logic
        
        info!(module_id = %module_id, "All dependencies met. Generating report for module.");

        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);
                generated_results.insert(module_id.clone(), format!("Error: Failed to create LLM client: {}", 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);
            }
        }
        
        // TODO: This is a placeholder. Implement proper financial data formatting and injection.
        context.insert("financial_data", "...");

        let prompt = match Tera::one_off(&module_config.prompt_template, &context, true) {
            Ok(p) => 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));
                continue;
            }
        };

        let content = match llm_client.generate_text(prompt).await {
            Ok(c) => c,
            Err(e) => {
                let err_msg = format!("LLM generation failed: {}", e);
                error!(module_id = %module_id, "{}", err_msg);
                generated_results.insert(module_id.clone(), format!("Error: {}", err_msg));
                continue;
            }
        };
        
        info!(module_id = %module_id, "Successfully generated content.");

        let result_to_persist = NewAnalysisResult {
            request_id: command.request_id,
            symbol: command.symbol.clone(),
            template_id: command.template_id.clone(),
            module_id: module_id.clone(),
            content: content.clone(),
            meta_data: serde_json::json!({ "model_id": module_config.model_id }),
        };

        if let Err(e) = persistence_client.create_analysis_result(result_to_persist).await {
            error!(module_id = %module_id, "Failed to persist analysis result: {}", e);
            // Decide if we should continue or fail the whole workflow
        }

        generated_results.insert(module_id.clone(), content);
    }

    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
        ))
    })?;

    Ok(LlmClient::new(
        provider.api_base_url.clone(),
        provider.api_key.clone().into(),
        module_config.model_id.clone(),
    ))
}

async fn fetch_data_and_configs(
    client: &PersistenceClient,
    symbol: &str,
) -> Result<(
    CompanyProfileDto,
    Vec<TimeSeriesFinancialDto>,
    LlmProvidersConfig,
    AnalysisTemplateSets,
)> {
    let (profile, financials, llm_providers, template_sets) = tokio::try_join!(
        client.get_company_profile(symbol),
        client.get_financial_statements(symbol),
        client.get_llm_providers_config(),
        client.get_analysis_template_sets(), // Changed from get_analysis_modules_config
    )?;
    Ok((profile, financials, llm_providers, template_sets))
}
