use crate::requirements::ProcessRequirements;
use cwl_core::{
    OneOrMany,
    documents::CWLDocument,
    inputs::{
        CommandInputParameter, CommandInputParameterType, CommandInputSchema, CommandInputType,
        InputSchema, InputType, WorkflowInputParameter,
    },
    outputs::{
        CommandOutputParameter, CommandOutputParameterType, CommandOutputType,
        ExpressionToolOutputParameter, OutputType, WorkflowOutputParameter,
    },
};
use std::collections::HashMap;

pub(crate) fn replace_schema_definitions(
    doc: &mut CWLDocument,
    requirements: &[ProcessRequirements],
) -> anyhow::Result<()> {
    let schema_defs = if let Some(sdr) = requirements
        .iter()
        .map(|i| {
            if let ProcessRequirements::SchemaDefRequirement(r) = i {
                Some(r)
            } else {
                None
            }
        })
        .next()
        .flatten()
    {
        Some(get_schema_definitions(&sdr.types)?)
    } else {
        None
    };

    //inputs are changed recursively... do we need that for outputs too??
    if let Some(defs) = schema_defs {
        match doc {
            CWLDocument::CommandLineTool(clt) => {
                add_schema_defs_to_command_inputs(&mut clt.inputs, &defs)?;
                add_schema_defs_to_command_outputs(&mut clt.outputs, &defs)?;
            }
            CWLDocument::ExpressionTool(et) => {
                add_schema_defs_to_inputs(&mut et.inputs, &defs)?;
                add_schema_defs_to_expression_outputs(&mut et.outputs, &defs)?;
            }
            CWLDocument::Workflow(wf) => {
                add_schema_defs_to_inputs(&mut wf.inputs, &defs)?;
                add_schema_defs_to_outputs(&mut wf.outputs, &defs)?;
            }
            CWLDocument::Operation(_) => {}
        }
    }

    Ok(())
}

fn add_schema_defs_to_command_inputs(
    inputs: &mut Vec<CommandInputParameter>,
    defs: &HashMap<String, serde_json::Value>,
) -> anyhow::Result<()> {
    for input in inputs {
        match &mut input.r#type {
            CommandInputParameterType::CommandInputType(OneOrMany::One(ty)) => {
                add_schema_defs_to_command_inputs_impl(ty, defs)?;
            }
            CommandInputParameterType::CommandInputType(OneOrMany::Many(tys)) => {
                for ty in tys {
                    add_schema_defs_to_command_inputs_impl(ty, defs)?;
                }
            }
            _ => {}
        }
    }
    Ok(())
}

fn add_schema_defs_to_command_inputs_impl(
    r#type: &mut CommandInputType,
    defs: &HashMap<String, serde_json::Value>,
) -> anyhow::Result<()> {
    match r#type {
        CommandInputType::CommandInputSchema(schema) => match &mut **schema {
            CommandInputSchema::Record(rec) => {
                if let Some(fields) = &mut rec.fields {
                    for field in fields {
                        match &mut field.r#type {
                            OneOrMany::One(item) => {
                                add_schema_defs_to_command_inputs_impl(item, defs)?;
                            }
                            OneOrMany::Many(items) => {
                                for item in items {
                                    add_schema_defs_to_command_inputs_impl(item, defs)?;
                                }
                            }
                        }
                    }
                }
            }
            CommandInputSchema::Array(arr) => match &mut arr.items {
                OneOrMany::One(item) => {
                    add_schema_defs_to_command_inputs_impl(item, defs)?;
                }
                OneOrMany::Many(items) => {
                    for item in items {
                        add_schema_defs_to_command_inputs_impl(item, defs)?;
                    }
                }
            },
            CommandInputSchema::Enum(_) => {}
        },
        CommandInputType::String(s) => {
            if let Some(def) = defs.get(&format!("#{s}")) {
                let new_type: CommandInputType = serde_json::from_value(def.clone())?;
                *r#type = new_type;
            } else if s.starts_with('#')
                && let Some(def) = defs.get(s)
            {
                let new_type: CommandInputType = serde_json::from_value(def.clone())?;
                *r#type = new_type;
            } else if s.contains('#')
                && let Some(ar) = s.split_once('#')
                && let Some(def) = defs.get(&format!("#{}", ar.1))
            {
                let new_type: CommandInputType = serde_json::from_value(def.clone())?;
                *r#type = new_type;
            }
        }
        CommandInputType::CWLType(_) => {}
    }
    Ok(())
}

fn add_schema_defs_to_inputs(
    inputs: &mut Vec<WorkflowInputParameter>,
    defs: &HashMap<String, serde_json::Value>,
) -> anyhow::Result<()> {
    for input in inputs {
        match &mut input.r#type {
            OneOrMany::One(ty) => add_schema_defs_to_inputs_impl(ty, defs)?,
            OneOrMany::Many(tys) => {
                for ty in tys {
                    add_schema_defs_to_inputs_impl(ty, defs)?;
                }
            }
        }
    }
    Ok(())
}

fn add_schema_defs_to_inputs_impl(
    r#type: &mut InputType,
    defs: &HashMap<String, serde_json::Value>,
) -> anyhow::Result<()> {
    match r#type {
        InputType::InputSchema(schema) => match &mut **schema {
            InputSchema::Record(rec) => {
                if let Some(fields) = &mut rec.fields {
                    for field in fields {
                        match &mut field.r#type {
                            OneOrMany::One(item) => {
                                add_schema_defs_to_inputs_impl(item, defs)?;
                            }
                            OneOrMany::Many(items) => {
                                for item in items {
                                    add_schema_defs_to_inputs_impl(item, defs)?;
                                }
                            }
                        }
                    }
                }
            }
            InputSchema::Array(arr) => match &mut arr.items {
                OneOrMany::One(item) => {
                    add_schema_defs_to_inputs_impl(item, defs)?;
                }
                OneOrMany::Many(items) => {
                    for item in items {
                        add_schema_defs_to_inputs_impl(item, defs)?;
                    }
                }
            },
            InputSchema::Enum(_) => {}
        },
        InputType::String(s) => {
            if let Some(def) = defs.get(&format!("#{s}")) {
                let new_type: InputType = serde_json::from_value(def.clone())?;
                *r#type = new_type;
            } else if s.starts_with('#')
                && let Some(def) = defs.get(s)
            {
                let new_type: InputType = serde_json::from_value(def.clone())?;
                *r#type = new_type;
            } else if s.contains('#')
                && let Some(ar) = s.split_once('#')
                && let Some(def) = defs.get(&format!("#{}", ar.1))
            {
                let new_type: InputType = serde_json::from_value(def.clone())?;
                *r#type = new_type;
            }
        }
        InputType::CWLType(_) => {}
    }

    Ok(())
}

fn add_schema_defs_to_command_outputs(
    inputs: &mut Vec<CommandOutputParameter>,
    defs: &HashMap<String, serde_json::Value>,
) -> anyhow::Result<()> {
    for input in inputs {
        if let CommandOutputParameterType::CommandOutputType(OneOrMany::One(
            CommandOutputType::String(s),
        )) = &mut input.r#type
            && let Some(def) = defs.get(&format!("#{s}"))
        {
            let new_type: CommandOutputParameterType = serde_json::from_value(def.clone())?;
            input.r#type = new_type;
        }
    }
    Ok(())
}

fn add_schema_defs_to_expression_outputs(
    inputs: &mut Vec<ExpressionToolOutputParameter>,
    defs: &HashMap<String, serde_json::Value>,
) -> anyhow::Result<()> {
    for input in inputs {
        if let OneOrMany::One(OutputType::String(s)) = &mut input.r#type
            && let Some(def) = defs.get(&format!("#{s}"))
        {
            let new_type: OneOrMany<OutputType> = serde_json::from_value(def.clone())?;
            input.r#type = new_type;
        }
    }
    Ok(())
}

fn add_schema_defs_to_outputs(
    inputs: &mut Vec<WorkflowOutputParameter>,
    defs: &HashMap<String, serde_json::Value>,
) -> anyhow::Result<()> {
    for input in inputs {
        if let CommandOutputParameterType::CommandOutputType(OneOrMany::One(
            CommandOutputType::String(s),
        )) = &mut input.r#type
            && let Some(def) = defs.get(&format!("#{s}"))
        {
            let new_type: CommandOutputParameterType = serde_json::from_value(def.clone())?;
            input.r#type = new_type;
        }
    }
    Ok(())
}

pub(crate) fn get_schema_definitions(
    value: &serde_json::Value,
) -> anyhow::Result<HashMap<String, serde_json::Value>> {
    let mut defs = extract_schema_definitions(value);

    if !defs.is_empty() {
        //replace recursive definitions
        let snapshot = &defs.clone();
        for item in defs.values_mut() {
            replace_schema_references(item, snapshot)?;
        }
    }
    Ok(defs)
}

fn extract_schema_definitions(value: &serde_json::Value) -> HashMap<String, serde_json::Value> {
    let mut schemas = HashMap::new();

    // Flatten: types can be a Sequence of Mappings, or a Sequence containing
    // a nested Sequence (when $import resolves a multi-type file)
    let type_defs: Vec<&serde_json::Value> = match value {
        serde_json::Value::Array(types) => types
            .iter()
            .flat_map(|item| match item {
                serde_json::Value::Array(inner) => inner.iter().collect::<Vec<_>>(),
                other => vec![other],
            })
            .collect(),
        mapping @ serde_json::Value::Object(_) => vec![mapping],
        _ => vec![],
    };

    for type_def in type_defs {
        if let serde_json::Value::Object(type_map) = type_def
            && let Some(serde_json::Value::String(name)) = type_map.get("name")
        {
            let mut type_def = type_def.clone();
            //validate names
            validate_field_name(&mut type_def);

            schemas.insert(format!("#{name}"), type_def);
        }
    }

    schemas
}

fn replace_schema_references(
    value: &mut serde_json::Value,
    schemas: &HashMap<String, serde_json::Value>,
) -> anyhow::Result<()> {
    match value {
        serde_json::Value::String(s) => {
            if s.starts_with('#') && schemas.contains_key(s) {
                *value = schemas[s].clone();
            } else if s.contains('#')
                && let Some(ar) = s.split_once('#')
            {
                //by doing this we accept that each id can be given only once!
                let s = format!("#{}", ar.1);
                if schemas.contains_key(&s) {
                    *value = schemas[&s].clone();
                }
            } else if schemas.contains_key(&format!("#{s}")) {
                *value = schemas[&format!("#{s}")].clone();
            }
        }
        serde_json::Value::Array(arr) => {
            for item in arr {
                replace_schema_references(item, schemas)?;
            }
        }
        serde_json::Value::Object(map) => {
            for (k, v) in map {
                if *k == serde_json::Value::String("name".to_string()) {
                    continue;
                }
                replace_schema_references(v, schemas)?;
            }
        }
        _ => {}
    }
    Ok(())
}

fn validate_field_name(schema: &mut serde_json::Value) {
    if let serde_json::Value::Object(record) = schema {
        //sanitize name
        if let Some(serde_json::Value::String(name)) = record.get_mut("name")
            && short_name(name) != name
        {
            *name = short_name(name).into();
        }

        if let Some(serde_json::Value::Array(fields)) = record.get_mut("fields") {
            for field in fields {
                validate_field_name(field);
            }
        }
    }
}

fn short_name(name: &str) -> &str {
    name.rsplit(['/', '#']).next().unwrap_or(name)
}
