use std::sync::Arc;

use roto::{List, Runtime};

use crate::context::{ExecutionContext, RunCancellation};
use crate::diagnostic::{Category, Diagnostic, LogEntry, RunDiagnostics, Severity};
use crate::error::{Error, Result};
use crate::loader::load_script;
use crate::module::ModuleRef;
use crate::request::RunRequest;
use crate::result::RunResult;

fn verify_registration(
    modules: &[ModuleRef],
    runtime: &mut Runtime<roto::Ctx<ExecutionContext>>,
) -> Result<()> {
    for module in modules {
        let spec = module.spec();
        for func in &spec.functions {
            let name = func.name.clone();
            let doc = func.doc.clone();
            let args: Vec<_> = func.args.iter().map(|s| s.as_str()).collect();

            let dummy_fn: fn(std::sync::Arc<str>) -> std::sync::Arc<str> = |v| v;
            let lib = match roto::Function::new(
                name.as_str(),
                doc.as_str(),
                args,
                dummy_fn,
                roto::location!(),
            ) {
                Ok(f) => {
                    let mut lib = roto::Library::new();
                    lib.add(roto::Item::Function(f));
                    lib
                }
                Err(e) => {
                    return Err(Error::runtime_contract(format!(
                        "spec function '{}' is invalid: {}",
                        func.name, e
                    ))
                    .with_primary_diagnostic("spec-function-invalid"));
                }
            };

            match runtime.add(lib) {
                Ok(()) => {
                    return Err(Error::runtime_contract(format!(
                        "module '{}' declared function '{}' in its spec, but did not register it",
                        spec.name, func.name
                    ))
                    .with_primary_diagnostic("spec-registration-mismatch"));
                }
                Err(e) => {
                    let err_str = e.to_string();
                    if !err_str.contains("declared twice") {
                        return Err(Error::internal(format!(
                            "unexpected registration error for '{}': {}",
                            func.name, err_str
                        ))
                        .with_primary_diagnostic("unexpected-registration-error"));
                    }
                }
            }
        }
    }
    Ok(())
}

pub struct Engine {
    modules: Vec<ModuleRef>,
}

impl Engine {
    pub fn new(modules: Vec<ModuleRef>) -> Self {
        Self { modules }
    }

    pub fn run_once(&self, request: &RunRequest) -> Result<RunResult> {
        self.run_once_with_cancellation(
            request,
            RunCancellation::with_timeout(request.limits.timeout_s),
        )
    }

    pub fn run_once_with_cancellation(
        &self,
        request: &RunRequest,
        cancellation: RunCancellation,
    ) -> Result<RunResult> {
        let mut collector = RunDiagnostics::default();

        request
            .validate()
            .map_err(|err| compose_error(err, &collector))?;
        cancellation
            .check_cancelled()
            .map_err(|err| compose_error(err, &collector))?;
        let loaded = load_script(&request.script, &request.limits).map_err(|err| {
            let category = err.kind().category();
            add_diagnostic_for(err, category, "load-script-failed")
        })?;
        cancellation
            .check_cancelled()
            .map_err(|err| compose_error(err, &collector))?;

        let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
            self.execute_loaded_script(&loaded.compile_path, request, cancellation, &mut collector)
        }));

        let cleanup_result = loaded.cleanup();

        for module in &self.modules {
            if let Err(e) = module.close() {
                collector.log(LogEntry {
                    severity: Severity::Error,
                    message: format!("module '{}' close failed: {}", module.spec().name, e),
                });
            }
        }

        match result {
            Ok(Ok(v)) => {
                cleanup_result
                    .map_err(|err| add_diagnostic_for(err, Category::Internal, "cleanup-failed"))?;
                Ok(v.with_diagnostics(collector))
            }
            Ok(Err(e)) => {
                cleanup_result
                    .map_err(|err| add_diagnostic_for(err, Category::Internal, "cleanup-failed"))?;
                Err(compose_error(e, &collector))
            }
            Err(e) => {
                if let Err(err) = cleanup_result {
                    collector.log(LogEntry {
                        severity: Severity::Error,
                        message: format!("cleanup failed after panic: {err}"),
                    });
                }
                let msg = if let Some(s) = e.downcast_ref::<&str>() {
                    s.to_string()
                } else if let Some(s) = e.downcast_ref::<String>() {
                    s.clone()
                } else {
                    "unknown panic".to_string()
                };
                let err = Error::panic(format!("script execution panicked: {msg}"));
                Err(compose_error(err, &collector))
            }
        }
    }

    fn execute_loaded_script(
        &self,
        compile_path: &std::path::Path,
        request: &RunRequest,
        cancellation: RunCancellation,
        collector: &mut RunDiagnostics,
    ) -> Result<RunResult> {
        cancellation.check_cancelled()?;
        let mut runtime = Runtime::new()
            .with_context_type::<ExecutionContext>()
            .map_err(|msg| {
                Error::internal(format!("register context type: {msg}"))
                    .with_primary_diagnostic("context-type-registration-failed")
            })?;
        cancellation.check_cancelled()?;

        for module in &self.modules {
            cancellation.check_cancelled()?;

            let module_prefix = format!("{}.", module.spec().name);
            let mut module_params = std::collections::HashMap::new();

            for (k, v) in &request.parameters {
                if let Some(stripped) = k.strip_prefix(&module_prefix) {
                    module_params.insert(stripped.to_string(), v.clone());
                } else if !k.contains('.') {
                    // Fallback for unprefixed parameters (e.g. from local CLI)
                    module_params.insert(k.clone(), v.clone());
                }
            }

            if let Err(err) = module.register(&mut runtime, &module_params) {
                collector.diagnostic(
                    Diagnostic::new(Category::Internal, Severity::Error, err.message_str())
                        .with_code("module-register-failed"),
                );
                return Err(err.with_primary_diagnostic("module-register-failed"));
            }
            cancellation.check_cancelled()?;
        }

        verify_registration(&self.modules, &mut runtime).map_err(|err| {
            let code = err
                .diagnostics()
                .first()
                .and_then(|d| d.code.as_deref())
                .unwrap_or("verify-registration-failed");
            collector.diagnostic(
                Diagnostic::new(err.kind().category(), Severity::Error, err.message_str())
                    .with_code(code),
            );
            compose_error(err, collector)
        })?;
        cancellation.check_cancelled()?;

        let mut pkg = match runtime.compile(compile_path) {
            Ok(pkg) => pkg,
            Err(report) => {
                let text = report.to_string();
                collector.diagnostic(
                    Diagnostic::new(Category::Compile, Severity::Error, &text)
                        .with_code("roto-compile-error"),
                );
                return Err(Error::compile(format!("Roto compilation failed\n{text}")));
            }
        };
        cancellation.check_cancelled()?;

        let mut ctx =
            ExecutionContext::from_limits_and_cancellation(&request.limits, cancellation.clone());
        ctx.check_cancelled()?;

        let output = match pkg.get_function::<fn() -> Arc<str>>(&request.entrypoint) {
            Ok(function) => {
                ctx.check_cancelled()?;
                if !request.args.is_empty() {
                    collector.diagnostic(
                        Diagnostic::new(
                            Category::RuntimeContract,
                            Severity::Error,
                            format!(
                                "entrypoint '{}' does not accept arguments",
                                request.entrypoint
                            ),
                        )
                        .with_code("entrypoint-arguments-mismatch"),
                    );
                    return Err(Error::runtime_contract(format!(
                        "entrypoint '{}' does not accept arguments",
                        request.entrypoint
                    )));
                }
                let output = function.call(&mut ctx);
                ctx.check_cancelled()?;
                output
            }
            Err(_no_arg_err) => {
                match pkg.get_function::<fn(List<Arc<str>>) -> Arc<str>>(&request.entrypoint) {
                    Ok(function) => {
                        ctx.check_cancelled()?;
                        let args = List::<Arc<str>>::new();
                        for arg in &request.args {
                            args.push(arg.clone().into());
                        }
                        let output = function.call(&mut ctx, args);
                        ctx.check_cancelled()?;
                        output
                    }
                    Err(_list_err) => {
                        collector.diagnostic(
                            Diagnostic::new(
                                Category::RuntimeContract,
                                Severity::Error,
                                format!(
                                    "entrypoint '{}' must match one of the supported signatures",
                                    request.entrypoint
                                ),
                            )
                            .with_code("entrypoint-signature-mismatch"),
                        );
                        return Err(Error::runtime_contract(format!(
                            "entrypoint '{}' must match one of the supported signatures: fn() -> String or fn(List<String>) -> String",
                            request.entrypoint
                        )));
                    }
                }
            }
        };

        let output_str = output.to_string();
        if output_str.len() > request.limits.max_output_bytes {
            collector.diagnostic(
                Diagnostic::new(
                    Category::RuntimeContract,
                    Severity::Error,
                    format!(
                        "output exceeds max_output_bytes: {} > {}",
                        output_str.len(),
                        request.limits.max_output_bytes
                    ),
                )
                .with_code("output-too-large"),
            );
            return Err(Error::resource_exhausted(format!(
                "output exceeds max_output_bytes: {} > {}",
                output_str.len(),
                request.limits.max_output_bytes
            )));
        }

        Ok(RunResult::new(output_str))
    }
}

fn compose_error(mut err: Error, collector: &RunDiagnostics) -> Error {
    for d in &collector.diagnostics {
        err = err.with_diagnostic(d.clone());
    }
    for l in &collector.logs {
        err = err.with_diagnostic(Diagnostic::new(Category::Internal, l.severity, &l.message));
    }
    err
}

fn add_diagnostic_for(err: Error, category: Category, code: &str) -> Error {
    let message = err.message_str().to_string();
    err.with_diagnostic(Diagnostic::new(category, Severity::Error, message).with_code(code))
}

#[cfg(test)]
mod tests {
    use std::sync::Arc;
    use std::time::{Duration, Instant};

    use roto::Runtime;

    use crate::context::ExecutionContext;
    use crate::diagnostic::{Category, Severity};
    use crate::engine::Engine;
    use crate::error::ErrorKind;
    use crate::module::{FunctionModule, FunctionSpec, Module, ModuleSpec};
    use crate::request::{RunLimits, RunRequest, ScriptSource};

    struct InvalidSpecModule;

    impl Module for InvalidSpecModule {
        fn spec(&self) -> ModuleSpec {
            ModuleSpec {
                name: String::from("broken"),
                doc: String::from("test"),
                capabilities: vec![],
                parameters: vec![],
                functions: vec![FunctionSpec {
                    name: String::new(),
                    doc: String::from("invalid"),
                    args: vec![],
                    return_type: String::from("String"),
                }],
            }
        }

        fn register(
            &self,
            _runtime: &mut Runtime<roto::Ctx<ExecutionContext>>,
            _params: &std::collections::HashMap<String, String>,
        ) -> crate::error::Result<()> {
            Ok(())
        }
    }

    fn test_request(source: &str) -> RunRequest {
        RunRequest {
            script: ScriptSource::Inline {
                source: source.into(),
            },
            entrypoint: "main".into(),
            args: vec![],
            parameters: std::collections::HashMap::new(),
            limits: RunLimits::default(),
        }
    }

    #[test]
    fn compile_error_includes_diagnostic() {
        let engine = Engine::new(vec![]);
        let req = test_request("fn broken {");
        let result = engine.run_once(&req);
        assert!(result.is_err());
        let err = result.unwrap_err();
        assert!(
            !err.diagnostics().is_empty(),
            "expected compile diagnostics"
        );
        let d = &err.diagnostics()[0];
        assert_eq!(d.category, Category::Compile);
        assert_eq!(d.code.as_deref(), Some("roto-compile-error"));
        assert_eq!(d.severity, Severity::Error);
        assert_eq!(err.kind(), ErrorKind::Compile);
    }

    #[test]
    fn signature_mismatch_produces_diagnostic() {
        let engine = Engine::new(vec![]);
        let req = test_request(r#"fn dummy() -> String { "hi" }"#);
        let result = engine.run_once(&req);
        assert!(result.is_err());
        let err = result.unwrap_err();
        assert!(
            err.diagnostics()
                .iter()
                .any(|d| d.code.as_deref() == Some("entrypoint-signature-mismatch")),
            "expected entrypoint-signature-mismatch diagnostic"
        );
        assert_eq!(err.kind(), ErrorKind::RuntimeContract);
    }

    #[test]
    fn output_too_large_produces_diagnostic() {
        let engine = Engine::new(vec![]);
        let req = RunRequest {
            script: ScriptSource::Inline {
                source: r#"fn main() -> String { "toolong" }"#.into(),
            },
            entrypoint: "main".into(),
            args: vec![],
            parameters: std::collections::HashMap::new(),
            limits: RunLimits {
                max_script_bytes: 64 * 1024,
                max_args: 64,
                max_arg_bytes: 64 * 1024,
                max_output_bytes: 4,
                timeout_s: 0,
                max_temp_bytes: 64 * 1024,
                metadata: None,
            },
        };
        let result = engine.run_once(&req);
        assert!(result.is_err());
        let err = result.unwrap_err();
        assert!(
            err.diagnostics()
                .iter()
                .any(|d| d.code.as_deref() == Some("output-too-large")),
            "expected output-too-large diagnostic"
        );
        assert_eq!(err.kind(), ErrorKind::ResourceExhausted);
    }

    #[test]
    fn success_result_can_carry_diagnostics() {
        let engine = Engine::new(vec![]);
        let req = test_request(r#"fn main() -> String { "hi" }"#);
        let result = engine.run_once(&req);
        assert!(result.is_ok());
        let run_result = result.unwrap();
        // Currently no diagnostics on success paths, but the field exists.
        assert!(
            run_result
                .diagnostics
                .map(|d: crate::diagnostic::RunDiagnostics| d.is_empty())
                .unwrap_or(true),
            "expected empty diagnostics on a successful run"
        );
    }

    #[test]
    fn zero_timeout_does_not_cancel_successful_run() {
        let engine = Engine::new(vec![]);
        let mut req = test_request(r#"fn main() -> String { "hi" }"#);
        req.limits.timeout_s = 0;

        let result = engine.run_once(&req).unwrap();

        assert_eq!(result.output, "hi");
    }

    #[test]
    fn pre_cancelled_run_returns_cancelled_error() {
        let engine = Engine::new(vec![]);
        let req = test_request(r#"fn main() -> String { "hi" }"#);
        let cancellation = crate::context::RunCancellation::none();
        cancellation.cancel();

        let err = engine
            .run_once_with_cancellation(&req, cancellation)
            .unwrap_err();

        assert_eq!(err.kind(), ErrorKind::Cancelled);
        assert!(
            err.diagnostics()
                .iter()
                .any(|d| d.code.as_deref() == Some("execution-cancelled")),
            "expected execution-cancelled diagnostic"
        );
    }

    #[test]
    fn expired_deadline_returns_timeout_diagnostic() {
        let engine = Engine::new(vec![]);
        let req = test_request(r#"fn main() -> String { "hi" }"#);
        let cancellation =
            crate::context::RunCancellation::new(Some(Instant::now() - Duration::from_millis(1)));

        let err = engine
            .run_once_with_cancellation(&req, cancellation)
            .unwrap_err();

        assert_eq!(err.kind(), ErrorKind::Cancelled);
        assert!(
            err.diagnostics()
                .iter()
                .any(|d| d.code.as_deref() == Some("timeout-exceeded")),
            "expected timeout-exceeded diagnostic"
        );
    }

    #[test]
    fn missing_spec_registration_produces_diagnostic() {
        let module = FunctionModule::new(
            ModuleSpec {
                name: String::from("missing"),
                doc: String::from("test"),
                capabilities: vec![],
                parameters: vec![],
                functions: vec![FunctionSpec::new(
                    "missing",
                    "echo",
                    "Declared but not registered.",
                    vec![String::from("value")],
                    "String",
                )],
            },
            |_, _| Ok(()),
            Arc::new(|| Ok(())),
        );
        let engine = Engine::new(vec![Arc::new(module)]);
        let req = test_request(r#"fn main() -> String { "ok" }"#);
        let result = engine.run_once(&req);
        assert!(result.is_err());
        let err = result.unwrap_err();
        assert!(
            err.diagnostics()
                .iter()
                .any(|d| d.code.as_deref() == Some("spec-registration-mismatch")),
            "expected spec-registration-mismatch diagnostic"
        );
        assert!(
            err.message_str().contains("did not register it"),
            "expected missing-registration message, got: {}",
            err.message_str()
        );
        assert_eq!(err.kind(), ErrorKind::RuntimeContract);
    }

    #[test]
    fn invalid_spec_registration_preserves_original_diagnostic_code() {
        let engine = Engine::new(vec![Arc::new(InvalidSpecModule)]);
        let req = test_request(r#"fn main() -> String { "ok" }"#);
        let err = engine.run_once(&req).unwrap_err();

        assert_eq!(err.kind(), ErrorKind::RuntimeContract);
        assert!(
            err.diagnostics()
                .iter()
                .any(|d| d.code.as_deref() == Some("spec-function-invalid")),
            "expected spec-function-invalid diagnostic"
        );
        assert!(
            !err.diagnostics()
                .iter()
                .any(|d| d.code.as_deref() == Some("spec-registration-mismatch")),
            "did not expect spec-registration-mismatch for invalid spec function"
        );
    }

    #[test]
    fn success_with_valid_module_and_script() {
        let module = crate::builtin::strings::module();
        let engine = Engine::new(vec![Arc::new(module)]);
        let req = test_request(r#"fn main() -> String { strings_uppercase("hi") }"#);
        let result = engine.run_once(&req);
        assert!(result.is_ok());
        assert_eq!(result.unwrap().output, "HI");
    }

    #[test]
    fn runtime_duplicate_registration_in_closure_fails() {
        let module_a = crate::builtin::strings::module();
        // module_b also tries to register strings_uppercase, which duplicates module_a.
        let module_b = FunctionModule::new(
            ModuleSpec {
                name: String::from("strings2"),
                doc: String::from("test"),
                capabilities: vec![],
                parameters: vec![],
                functions: vec![FunctionSpec::new(
                    "strings2",
                    "uppercase",
                    "doc",
                    vec![String::from("value")],
                    "String",
                )],
            },
            |runtime: &mut Runtime<roto::Ctx<ExecutionContext>>,
             _params: &std::collections::HashMap<String, String>| {
                crate::module::register_fn(
                    runtime,
                    "strings_uppercase",
                    "dup",
                    &["value"],
                    |v: Arc<str>| -> Arc<str> { v },
                )
            },
            Arc::new(|| Ok(())),
        );
        let engine = Engine::new(vec![Arc::new(module_a), Arc::new(module_b)]);
        let req = test_request(r#"fn main() -> String { "ok" }"#);
        let result = engine.run_once(&req);
        assert!(result.is_err());
        let err = result.unwrap_err();
        assert!(
            err.diagnostics()
                .iter()
                .any(|d| d.code.as_deref() == Some("module-register-failed")),
            "expected module-register-failed diagnostic"
        );
        assert_eq!(err.kind(), ErrorKind::RuntimeContract);
    }

    #[test]
    fn entrypoint_arguments_mismatch_produces_diagnostic() {
        let engine = Engine::new(vec![]);
        let req = RunRequest {
            script: ScriptSource::Inline {
                source: r#"fn main() -> String { "hi" }"#.into(),
            },
            entrypoint: "main".into(),
            args: vec![String::from("unexpected")],
            parameters: std::collections::HashMap::new(),
            limits: RunLimits::default(),
        };
        let result = engine.run_once(&req);
        assert!(result.is_err());
        let err = result.unwrap_err();
        assert!(
            err.diagnostics()
                .iter()
                .any(|d| d.code.as_deref() == Some("entrypoint-arguments-mismatch")),
            "expected entrypoint-arguments-mismatch diagnostic"
        );
        assert_eq!(err.kind(), ErrorKind::RuntimeContract);
    }

    #[test]
    fn module_parameters_are_filtered_and_prefixes_stripped() {
        use std::sync::atomic::{AtomicBool, Ordering};

        let called_correctly = Arc::new(AtomicBool::new(false));
        let called_clone = called_correctly.clone();

        let module = crate::module::FunctionModule::new(
            crate::module::ModuleSpec {
                name: "test_mod".into(),
                doc: "".into(),
                capabilities: vec![],
                parameters: vec![],
                functions: vec![],
            },
            move |_rt, params| {
                if params.get("API_KEY").map(String::as_str) == Some("secret123")
                    && params.get("unprefixed_key").map(String::as_str) == Some("fallback")
                    && params.get("other_mod.API_KEY").is_none()
                {
                    called_clone.store(true, Ordering::SeqCst);
                }
                Ok(())
            },
            Arc::new(|| Ok(())),
        );

        let engine = Engine::new(vec![Arc::new(module)]);
        let mut parameters = std::collections::HashMap::new();
        parameters.insert("test_mod.API_KEY".into(), "secret123".into());
        parameters.insert("other_mod.API_KEY".into(), "wrong".into());
        parameters.insert("unprefixed_key".into(), "fallback".into());

        let req = RunRequest {
            script: ScriptSource::Inline {
                source: r#"fn main() -> String { "ok" }"#.into(),
            },
            entrypoint: "main".into(),
            args: vec![],
            parameters,
            limits: RunLimits::default(),
        };

        let result = engine.run_once(&req);
        assert!(result.is_ok());
        assert!(called_correctly.load(Ordering::SeqCst), "module did not receive expected parameters");
    }
}
