use std::{
    collections::HashMap,
    path::PathBuf,
};

use crate::{
    collection::Collection,
    request::{RequestTemplate, VariableStore, expand_templates},
    schema::SchemaRegistry,
};

use super::{
    Rule, context,
    declarations::{
        DeclarationSpanIndex, parse_request_collection, register_declaration,
        remember_declaration_span, validate_schema_registry,
    },
    legacy_header::normalize_legacy_collection_header,
    preprocessor,
    protocol::SessionRequestTarget,
    request::{
        parse_request_template,
    },
    spans::{
        prompt_error_to_span, template_error_to_collection_error,
        template_error_to_span,
    },
    variables::assign_global_variable,
};

pub fn parse_collection(
    input: &str,
    working_dir: PathBuf,
) -> Result<Collection, pest::error::Error<Rule>> {
    let preprocessed = preprocessor::preprocess(input, &working_dir).map_err(|e| {
        pest::error::Error::new_from_span(
            pest::error::ErrorVariant::CustomError {
                message: e.to_string(),
            },
            pest::Span::new(input, 0, input.len()).unwrap(),
        )
    })?;
    let preprocessed = normalize_legacy_collection_header(preprocessed.as_str());

    log::debug!("PREPROCESSED COMPLETE:\n{}", preprocessed);

    let collection = parse_request_collection(preprocessed.as_str())?;

    let mut name = String::new();
    let mut description = String::new();
    let mut request_templates: Vec<RequestTemplate> = Vec::new();
    let mut global_store = VariableStore::new();
    let mut global_headers: HashMap<String, String> = HashMap::new();
    let mut global_queries: HashMap<String, String> = HashMap::new();
    let mut global_callbacks: Vec<String> = vec![];
    let mut schema_registry = SchemaRegistry::default();
    let mut declaration_spans = DeclarationSpanIndex::new();
    let mut session_targets: HashMap<String, SessionRequestTarget> = HashMap::new();

    for pair in collection.into_inner() {
        match pair.as_rule() {
            Rule::collection_name => {
                name = pair.as_str().trim().to_string();
            }

            Rule::collection_description => {
                description.push_str(pair.as_str().trim());
            }

            Rule::variable => {
                let span = pair.as_span();
                let mut inner_pairs = pair.into_inner();
                let key = inner_pairs.next().unwrap().as_str().trim().to_string();
                let value = inner_pairs.next().unwrap().as_str().to_string();

                assign_global_variable(&mut global_store, key, value.as_str(), &working_dir)
                    .map_err(|err| template_error_to_span(err, span.clone()))?;
            }

            Rule::header => {
                let span = pair.as_span();
                let mut inner_pairs = pair.into_inner();
                let key = inner_pairs.next().unwrap().as_str().trim().to_string();
                let value = inner_pairs.next().unwrap().as_str().trim().to_string();

                let value = context::try_inject_from_prompt(&value)
                    .map_err(|err| prompt_error_to_span(err, span.clone()))?;
                global_headers.insert(key, value);
            }

            Rule::query => {
                let span = pair.as_span();
                let mut inner_pairs = pair.into_inner();
                let key = inner_pairs.next().unwrap().as_str().trim().to_string();
                let value = inner_pairs.next().unwrap().as_str().trim().to_string();

                let value = context::try_inject_from_prompt(&value)
                    .map_err(|err| prompt_error_to_span(err, span.clone()))?;
                global_queries.insert(key, value);
            }

            Rule::callback => {
                global_callbacks.push(pair.as_str().strip_prefix('!').unwrap().to_string());
            }

            Rule::declaration => {
                remember_declaration_span(&mut declaration_spans, &pair);
                register_declaration(&mut schema_registry, pair)?;
            }

            Rule::requests => {
                validate_schema_registry(
                    &schema_registry,
                    preprocessed.as_str(),
                    &declaration_spans,
                )?;
                for request_pair in pair.into_inner() {
                    let template = parse_request_template(
                        request_pair,
                        &global_store,
                        &global_headers,
                        &global_queries,
                        &global_callbacks,
                        &schema_registry,
                        &working_dir,
                        &mut session_targets,
                    )?;
                    request_templates.push(template);
                }
            }

            Rule::EOI => {}

            _ => {
                unreachable!("unexpected rule: {:?}", pair.as_rule());
            }
        }
    }

    let requests = expand_templates(request_templates)
        .map_err(|err| template_error_to_collection_error(err, preprocessed.as_str()))?;

    Ok(Collection {
        name,
        description,
        schema_registry,
        requests,
    })
}