use std::collections::BTreeMap;

use serde::{Deserialize, Serialize};

use crate::adapter::{AdapterPath, AdapterRegistry, PlanningPolicy};
use crate::alignment::{alignment_metadata, alignment_mode_from_fusion};
use crate::error::{DataError, Result};
use crate::ids::{RepresentationId, SourceId};
use crate::model::{DatasetSchema, SourceDescriptor};
use crate::plan::{DataPlan, DataPlanStep, DataPlanStepKind, FitScope};
use crate::ModelInputSpec;

#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
pub struct DataPlanRequest {
    pub id: String,
    #[serde(default)]
    pub source_ids: Option<Vec<SourceId>>,
    #[serde(default)]
    pub planning_policy: PlanningPolicy,
}

impl DataPlanRequest {
    pub fn new(id: impl Into<String>) -> Self {
        Self {
            id: id.into(),
            source_ids: None,
            planning_policy: PlanningPolicy::default(),
        }
    }
}

struct ResolvedSource<'a> {
    source: &'a SourceDescriptor,
    target_representation: RepresentationId,
    path: AdapterPath,
}

pub fn plan_model_input(
    schema: &DatasetSchema,
    model_input: &ModelInputSpec,
    adapters: &AdapterRegistry,
    request: &DataPlanRequest,
) -> Result<DataPlan> {
    schema.validate()?;
    model_input.validate()?;
    validate_request(request)?;

    let candidate_sources = candidate_sources(schema, request)?;
    let mut steps = Vec::new();
    let mut output_representation = None;
    let mut adapt_idx = 0usize;
    let mut align_idx = 0usize;

    for port in &model_input.ports {
        let resolved = candidate_sources
            .iter()
            .filter_map(|source| resolve_source(source, port, adapters, &request.planning_policy))
            .collect::<Vec<_>>();

        if resolved.is_empty() {
            if port.optional {
                continue;
            }
            return Err(DataError::Validation(format!(
                "no source can satisfy required model input port `{}`",
                port.name
            )));
        }
        if resolved.len() > 1 && !port.multi_source {
            return Err(DataError::Validation(format!(
                "model input port `{}` accepts one source but {} sources resolved",
                port.name,
                resolved.len()
            )));
        }

        let port_output_representation = resolved[0].target_representation.clone();
        let mut join_inputs = Vec::new();
        for source_plan in resolved {
            let source_output = source_output_id(&source_plan.source.id);
            steps.push(DataPlanStep {
                kind: DataPlanStepKind::Materialize,
                source_id: Some(source_plan.source.id.clone()),
                adapter_id: None,
                input_representation: None,
                output_representation: Some(source_plan.source.native_representation.id.clone()),
                fit_scope: FitScope::Stateless,
                requires_user_choice: false,
                metadata: BTreeMap::from([(
                    "output".to_string(),
                    serde_json::Value::String(source_output.clone()),
                )]),
            });

            let mut current_output = source_output;
            let mut current_representation = source_plan.source.native_representation.id.clone();
            for adapter in source_plan.path.adapters {
                let step_output = format!("step:adapt:{adapt_idx}");
                steps.push(DataPlanStep {
                    kind: DataPlanStepKind::Adapt,
                    source_id: Some(source_plan.source.id.clone()),
                    adapter_id: Some(adapter.id),
                    input_representation: Some(current_representation),
                    output_representation: Some(adapter.output_representation.clone()),
                    fit_scope: adapter.fit_scope,
                    requires_user_choice: false,
                    metadata: BTreeMap::from([
                        (
                            "input".to_string(),
                            serde_json::Value::String(current_output),
                        ),
                        (
                            "output".to_string(),
                            serde_json::Value::String(step_output.clone()),
                        ),
                    ]),
                });
                current_output = step_output;
                current_representation = adapter.output_representation;
                adapt_idx += 1;
            }
            join_inputs.push(current_output);
        }

        let join_inputs = if join_inputs.len() > 1 {
            let align_output = format!("step:align:{align_idx}");
            let align_input_representation = output_representation
                .clone()
                .unwrap_or_else(|| resolved_target_representation(&steps));
            let alignment_mode = alignment_mode_from_fusion(model_input.default_fusion.as_ref())?;
            steps.push(DataPlanStep {
                kind: DataPlanStepKind::Align,
                source_id: None,
                adapter_id: None,
                input_representation: Some(align_input_representation),
                output_representation: Some(port_output_representation.clone()),
                fit_scope: FitScope::Stateless,
                requires_user_choice: false,
                metadata: alignment_metadata(join_inputs, align_output.clone(), alignment_mode)?,
            });
            align_idx += 1;
            vec![align_output]
        } else {
            join_inputs
        };

        steps.push(DataPlanStep {
            kind: DataPlanStepKind::Join,
            source_id: None,
            adapter_id: None,
            input_representation: Some(
                output_representation
                    .clone()
                    .unwrap_or_else(|| resolved_target_representation(&steps)),
            ),
            output_representation: Some(port_output_representation.clone()),
            fit_scope: FitScope::Stateless,
            requires_user_choice: false,
            metadata: BTreeMap::from([
                (
                    "inputs".to_string(),
                    serde_json::Value::Array(
                        join_inputs
                            .into_iter()
                            .map(serde_json::Value::String)
                            .collect(),
                    ),
                ),
                (
                    "output".to_string(),
                    serde_json::Value::String(format!("port:{}", port.name)),
                ),
            ]),
        });
        output_representation = Some(port_output_representation);
    }

    let output_representation = output_representation
        .ok_or_else(|| DataError::Validation("model input plan produced no outputs".to_string()))?;
    let plan = DataPlan {
        id: request.id.clone(),
        steps,
        output_representation,
        issues: Vec::new(),
    };
    plan.validate()?;
    Ok(plan)
}

fn validate_request(request: &DataPlanRequest) -> Result<()> {
    if request.id.trim().is_empty() {
        return Err(DataError::Validation(
            "data plan request id is empty".to_string(),
        ));
    }
    if let Some(source_ids) = &request.source_ids {
        let mut seen = std::collections::BTreeSet::new();
        for source_id in source_ids {
            if !seen.insert(source_id) {
                return Err(DataError::Validation(format!(
                    "data plan request contains duplicate source `{source_id}`"
                )));
            }
        }
    }
    Ok(())
}

fn candidate_sources<'a>(
    schema: &'a DatasetSchema,
    request: &DataPlanRequest,
) -> Result<Vec<&'a SourceDescriptor>> {
    let mut sources = schema.sources.iter().collect::<Vec<_>>();
    sources.sort_by(|left, right| left.id.cmp(&right.id));
    if let Some(requested) = &request.source_ids {
        let by_id = sources
            .iter()
            .map(|source| (&source.id, *source))
            .collect::<BTreeMap<_, _>>();
        requested
            .iter()
            .map(|source_id| {
                by_id.get(source_id).copied().ok_or_else(|| {
                    DataError::Validation(format!(
                        "data plan request references unknown source `{source_id}`"
                    ))
                })
            })
            .collect()
    } else {
        Ok(sources)
    }
}

fn resolve_source<'a>(
    source: &'a SourceDescriptor,
    port: &crate::InputPortSpec,
    adapters: &AdapterRegistry,
    policy: &PlanningPolicy,
) -> Option<ResolvedSource<'a>> {
    for target_type in &port.accepted_types {
        for target_representation in &port.accepted_representations {
            let resolution = adapters.find_path(
                &source.type_id,
                &source.native_representation.id,
                target_type,
                target_representation,
                policy,
            );
            if let Some(path) = resolution.path {
                return Some(ResolvedSource {
                    source,
                    target_representation: target_representation.clone(),
                    path,
                });
            }
        }
    }
    None
}

fn source_output_id(source_id: &SourceId) -> String {
    format!("src:{source_id}")
}

fn resolved_target_representation(steps: &[DataPlanStep]) -> RepresentationId {
    steps
        .iter()
        .rev()
        .find_map(|step| step.output_representation.clone())
        .expect("at least one source step was emitted")
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::adapter::AdapterRegistrySpec;
    use crate::ids::{RepresentationId, SourceId, TypeId};
    use crate::plan::DataPlan;

    fn load_schema() -> DatasetSchema {
        serde_json::from_str(include_str!(
            "../../../examples/fixtures/oof_campaign/schema_nir_6_samples.json"
        ))
        .unwrap()
    }

    fn load_model_input() -> ModelInputSpec {
        serde_json::from_str(include_str!(
            "../../../examples/fixtures/oof_campaign/model_input_tabular_numeric.json"
        ))
        .unwrap()
    }

    fn load_registry() -> AdapterRegistry {
        let spec: AdapterRegistrySpec = serde_json::from_str(include_str!(
            "../../../examples/fixtures/oof_campaign/adapter_registry_signal_to_tabular.json"
        ))
        .unwrap();
        AdapterRegistry::from_spec(spec).unwrap()
    }

    #[test]
    fn fixture_plan_matches_expected_json() {
        let plan = plan_model_input(
            &load_schema(),
            &load_model_input(),
            &load_registry(),
            &DataPlanRequest {
                id: "nir-to-tabular".to_string(),
                source_ids: Some(vec![SourceId::new("nir").unwrap()]),
                planning_policy: PlanningPolicy::default(),
            },
        )
        .unwrap();
        let expected: DataPlan = serde_json::from_str(include_str!(
            "../../../examples/fixtures/oof_campaign/expected_data_plan_nir_to_tabular.json"
        ))
        .unwrap();

        assert_eq!(plan, expected);
    }

    #[test]
    fn planner_refuses_missing_required_port_source() {
        let err = plan_model_input(
            &load_schema(),
            &load_model_input(),
            &load_registry(),
            &DataPlanRequest {
                id: "missing-source".to_string(),
                source_ids: Some(vec![SourceId::new("other").unwrap()]),
                planning_policy: PlanningPolicy::default(),
            },
        )
        .unwrap_err();

        assert!(err.to_string().contains("unknown source"));
    }

    #[test]
    fn planner_emits_alignment_step_for_multi_source_port() {
        let mut schema = load_schema();
        let mut chem = schema.sources[0].clone();
        chem.id = SourceId::new("chem").unwrap();
        chem.name = "Chemistry".to_string();
        chem.type_id = TypeId::new("table").unwrap();
        chem.native_representation.id = RepresentationId::new("tabular_numeric").unwrap();
        chem.native_representation.type_id = TypeId::new("table").unwrap();
        chem.native_representation.container = "dataframe".to_string();
        schema.sources.push(chem);

        let plan = plan_model_input(
            &schema,
            &load_model_input(),
            &load_registry(),
            &DataPlanRequest {
                id: "nir-chem-to-tabular".to_string(),
                source_ids: Some(vec![
                    SourceId::new("nir").unwrap(),
                    SourceId::new("chem").unwrap(),
                ]),
                planning_policy: PlanningPolicy::default(),
            },
        )
        .unwrap();

        let align = plan
            .steps
            .iter()
            .find(|step| step.kind == DataPlanStepKind::Align)
            .expect("multi-source plan should expose an Align step");
        assert_eq!(align.metadata["alignment"], serde_json::json!("left"));
        assert_eq!(align.metadata["output"], serde_json::json!("step:align:0"));
        assert_eq!(
            align.metadata["inputs"].as_array().unwrap().len(),
            2,
            "alignment should receive both source feature streams"
        );

        let join = plan
            .steps
            .iter()
            .find(|step| step.kind == DataPlanStepKind::Join)
            .unwrap();
        assert_eq!(join.metadata["inputs"], serde_json::json!(["step:align:0"]));
    }
}
