use std::net::SocketAddr;
use std::sync::Arc;

use tokio::signal;
use tokio_util::sync::CancellationToken;
use tonic::{Request, Response, Status};

use crate::context::RunCancellation;
use crate::error::{Error, Result};
use crate::service::AgentService;

// Include the generated protobuf code.
pub mod pb {
    tonic::include_proto!("sandlet.v1");
}

use pb::agent_service_server::{AgentService as PbAgentService, AgentServiceServer};
use pb::*;

fn map_error_kind(kind: crate::error::ErrorKind) -> i32 {
    use crate::error::ErrorKind;
    match kind {
        ErrorKind::InvalidRequest => pb::ErrorKind::InvalidRequest as i32,
        ErrorKind::Compile => pb::ErrorKind::Compile as i32,
        ErrorKind::RuntimeContract => pb::ErrorKind::RuntimeContract as i32,
        ErrorKind::RuntimeExecution => pb::ErrorKind::RuntimeExecution as i32,
        ErrorKind::Transport => pb::ErrorKind::Transport as i32,
        ErrorKind::ResourceExhausted => pb::ErrorKind::ResourceExhausted as i32,
        ErrorKind::Cancelled => pb::ErrorKind::Cancelled as i32,
        ErrorKind::Internal => pb::ErrorKind::Internal as i32,
        ErrorKind::Panic => pb::ErrorKind::Panic as i32,
    }
}

fn map_severity(s: crate::diagnostic::Severity) -> i32 {
    use crate::diagnostic::Severity;
    match s {
        Severity::Error => diagnostic::Severity::Error as i32,
        Severity::Warning => diagnostic::Severity::Warning as i32,
        Severity::Info => diagnostic::Severity::Info as i32,
        Severity::Debug => diagnostic::Severity::Debug as i32,
    }
}

fn map_category(c: crate::diagnostic::Category) -> i32 {
    use crate::diagnostic::Category;
    match c {
        Category::InvalidRequest => diagnostic::Category::InvalidRequest as i32,
        Category::Compile => diagnostic::Category::Compile as i32,
        Category::RuntimeContract => diagnostic::Category::RuntimeContract as i32,
        Category::RuntimeExecution => diagnostic::Category::RuntimeExecution as i32,
        Category::Transport => diagnostic::Category::Transport as i32,
        Category::ResourceExhausted => diagnostic::Category::ResourceExhausted as i32,
        Category::Cancelled => diagnostic::Category::Cancelled as i32,
        Category::Internal => diagnostic::Category::Internal as i32,
        Category::Panic => diagnostic::Category::Panic as i32,
    }
}

fn map_diagnostic(d: &crate::diagnostic::Diagnostic) -> pb::Diagnostic {
    pb::Diagnostic {
        category: map_category(d.category),
        severity: map_severity(d.severity),
        message: d.message.clone(),
        code: d.code.clone().unwrap_or_default(),
        location: d.location.as_ref().map(|l| pb::SourceLocation {
            file: l.file.clone(),
            line: l.line.unwrap_or(0),
            column: l.column.unwrap_or(0),
            start_byte: l.start_byte.unwrap_or(0) as u64,
            end_byte: l.end_byte.unwrap_or(0) as u64,
        }),
    }
}

fn map_log_entry(l: &crate::diagnostic::LogEntry) -> pb::LogEntry {
    pb::LogEntry {
        severity: map_severity(l.severity),
        message: l.message.clone(),
    }
}

pub struct GrpcServer {
    service: Arc<AgentService>,
}

struct RpcCancellationGuard {
    cancellation: RunCancellation,
    armed: bool,
}

impl RpcCancellationGuard {
    fn new(cancellation: RunCancellation) -> Self {
        Self {
            cancellation,
            armed: true,
        }
    }

    fn disarm(&mut self) {
        self.armed = false;
    }
}

impl Drop for RpcCancellationGuard {
    fn drop(&mut self) {
        if self.armed {
            self.cancellation.cancel();
        }
    }
}

impl GrpcServer {
    pub fn new(service: Arc<AgentService>) -> Self {
        Self { service }
    }

    pub fn serve(&self, listen_addr: &str) -> Result<()> {
        let addr: SocketAddr = listen_addr
            .parse()
            .map_err(|err| Error::invalid_request(format!("invalid listen address: {}", err)))?;

        let rt = tokio::runtime::Runtime::new()
            .map_err(|err| Error::transport(format!("failed to start tokio runtime: {}", err)))?;

        rt.block_on(async move {
            let server = AgentServiceServer::new(GrpcServer {
                service: Arc::clone(&self.service),
            });

            let token = CancellationToken::new();
            let token_clone = token.clone();
            tokio::spawn(async move {
                let _ = signal::ctrl_c().await;
                eprintln!("received shutdown signal");
                token_clone.cancel();
            });

            for m in &self.service.describe().modules {
                eprintln!("enabling module {}", m.name);
            }
            eprintln!("gRPC server listening on {}", addr);

            tonic::transport::Server::builder()
                .add_service(server)
                .serve_with_shutdown(addr, token.cancelled())
                .await
                .map_err(|err| Error::transport(format!("gRPC server error: {}", err)))?;

            Ok(())
        })
    }
}

fn map_run_success(result: crate::result::RunResult) -> RunResponse {
    let (diagnostics, logs) = if let Some(d) = result.diagnostics {
        (
            d.diagnostics.iter().map(map_diagnostic).collect(),
            d.logs.iter().map(map_log_entry).collect(),
        )
    } else {
        (Vec::new(), Vec::new())
    };

    RunResponse {
        result: Some(run_response::Result::Success(RunSuccess {
            output: result.output,
            diagnostics,
            logs,
        })),
    }
}

fn map_run_error(err: Error) -> RunResponse {
    let kind_str = err.kind().code_str();
    let diagnostics: Vec<pb::Diagnostic> = err.diagnostics().iter().map(map_diagnostic).collect();
    RunResponse {
        result: Some(run_response::Result::Failure(RunFailure {
            message: format!("{}: {}", kind_str, err.message_str()),
            diagnostics,
            kind: map_error_kind(err.kind()),
        })),
    }
}

fn map_limits(limits: Option<pb::RunLimits>) -> crate::request::RunLimits {
    let defaults = crate::request::RunLimits::default();
    match limits {
        None => defaults,
        Some(limits) => crate::request::RunLimits {
            max_script_bytes: limits
                .max_script_bytes
                .map(|value| value as usize)
                .unwrap_or(defaults.max_script_bytes),
            max_args: limits
                .max_args
                .map(|value| value as usize)
                .unwrap_or(defaults.max_args),
            max_arg_bytes: limits
                .max_arg_bytes
                .map(|value| value as usize)
                .unwrap_or(defaults.max_arg_bytes),
            max_output_bytes: limits
                .max_output_bytes
                .map(|value| value as usize)
                .unwrap_or(defaults.max_output_bytes),
            timeout_s: limits.timeout_s.unwrap_or(defaults.timeout_s),
            max_temp_bytes: limits
                .max_temp_bytes
                .map(|value| value as usize)
                .unwrap_or(defaults.max_temp_bytes),
            metadata: limits.metadata,
        },
    }
}

#[tonic::async_trait]
impl PbAgentService for GrpcServer {
    async fn run_once(
        &self,
        request: Request<RunRequest>,
    ) -> std::result::Result<Response<RunResponse>, Status> {
        let req = request.into_inner();
        let run_request = crate::request::RunRequest {
            script: match req.script.as_str() {
                "" => crate::request::ScriptSource::Inline {
                    source: String::new(),
                },
                source => crate::request::ScriptSource::Inline {
                    source: source.to_string(),
                },
            },
            entrypoint: String::from("main"),
            args: req.args,
            parameters: req.parameters.into_iter().collect(),
            limits: map_limits(req.limits),
        };
        let cancellation = RunCancellation::with_timeout(run_request.limits.timeout_s);
        let mut request_guard = RpcCancellationGuard::new(cancellation.clone());
        let service = Arc::clone(&self.service);
        let run_task = tokio::task::spawn_blocking(move || {
            service.run_once_with_cancellation(&run_request, cancellation)
        });

        let run_result = run_task.await;
        request_guard.disarm();

        let response = match run_result {
            Ok(Ok(result)) => map_run_success(result),
            Ok(Err(err)) => map_run_error(err),
            Err(join_err) if join_err.is_cancelled() => map_run_error(
                Error::cancelled("gRPC run task cancelled")
                    .with_primary_diagnostic("execution-cancelled"),
            ),
            Err(join_err) => map_run_error(
                Error::internal(format!("gRPC run task failed: {join_err}"))
                    .with_primary_diagnostic("runtime-task-join-failed"),
            ),
        };
        Ok(Response::new(response))
    }

    async fn describe(
        &self,
        _request: Request<pb::DescribeRequest>,
    ) -> std::result::Result<Response<DescribeResponse>, Status> {
        let metadata = self.service.describe();
        let modules: Vec<pb::Module> = metadata
            .modules
            .iter()
            .map(|m| pb::Module {
                name: m.name.clone(),
                description: m.doc.clone(),
                capabilities: m.capabilities.clone(),
                parameters: m
                    .parameters
                    .iter()
                    .map(|p| pb::ModuleParameter {
                        name: p.name.clone(),
                        description: p.doc.clone(),
                        secret: p.secret,
                    })
                    .collect(),
                functions: m
                    .functions
                    .iter()
                    .map(|f| pb::Function {
                        name: f.name.clone(),
                        description: f.doc.clone(),
                        arg_names: f.args.clone(),
                        return_type_hint: f.return_type.clone(),
                    })
                    .collect(),
            })
            .collect();

        Ok(Response::new(DescribeResponse {
            agent_name: metadata.agent_name.clone(),
            sdk_version: metadata.sdk_version.clone(),
            modules,
            declared_limits: metadata.declared_limits.clone(),
            metadata: metadata.metadata.clone(),
        }))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::result::AgentMetadata;
    use crate::service::AgentService;

    fn make_service() -> Arc<AgentService> {
        let engine = Arc::new(crate::engine::Engine::new(vec![]));
        let metadata = AgentMetadata {
            agent_name: "test-agent".into(),
            modules: vec![],
            sdk_version: "0.1.0".to_string(),
            declared_limits: std::collections::HashMap::new(),
            metadata: std::collections::HashMap::new(),
        };
        Arc::new(AgentService::new(engine, metadata))
    }

    #[test]
    fn map_limits_handles_none_and_partial_values() {
        let none_limits = super::map_limits(None);
        assert_eq!(none_limits.max_script_bytes, 64 * 1024);
        assert!(none_limits.metadata.is_none());

        let some_empty = super::map_limits(Some(RunLimits {
            max_script_bytes: Some(64 * 1024),
            max_args: None,
            max_arg_bytes: None,
            max_output_bytes: None,
            timeout_s: None,
            max_temp_bytes: None,
            metadata: Some(String::new()),
        }));
        assert_eq!(some_empty.metadata, Some(String::new()));

        let some_value = super::map_limits(Some(RunLimits {
            max_script_bytes: Some(64 * 1024),
            max_args: None,
            max_arg_bytes: None,
            max_output_bytes: None,
            timeout_s: None,
            max_temp_bytes: None,
            metadata: Some("trace-id-123".into()),
        }));
        assert_eq!(some_value.metadata, Some("trace-id-123".into()));
    }

    #[tokio::test]
    async fn run_once_success_includes_diagnostics_when_present() {
        let engine = Arc::new(crate::engine::Engine::new(vec![]));
        let metadata = AgentMetadata {
            agent_name: "diag-agent".into(),
            modules: vec![],
            sdk_version: "0.1.0".to_string(),
            declared_limits: std::collections::HashMap::new(),
            metadata: std::collections::HashMap::new(),
        };
        let service = Arc::new(AgentService::new(engine, metadata));
        let server = GrpcServer::new(Arc::clone(&service));

        let request = tonic::Request::new(RunRequest {
            script: r#"fn main() -> String { "hello" }"#.to_string(),
            args: vec![],
            parameters: std::collections::HashMap::new(),
            limits: None,
        });

        let response = PbAgentService::run_once(&server, request)
            .await
            .unwrap()
            .into_inner();
        let success = match response.result.expect("a result") {
            super::pb::run_response::Result::Success(s) => s,
            super::pb::run_response::Result::Failure(f) => {
                panic!("unexpected failure: {}", f.message)
            }
        };
        assert_eq!(success.output, "hello");
        // Currently no diagnostics on success, but the field exists.
        assert!(success.diagnostics.is_empty());
    }

    #[tokio::test]
    async fn run_once_failure_maps_structured_diagnostics() {
        let server = GrpcServer::new(make_service());

        let request = tonic::Request::new(RunRequest {
            script: "fn broken { ".to_string(),
            args: vec![],
            parameters: std::collections::HashMap::new(),
            limits: Some(RunLimits {
                max_script_bytes: Some(64 * 1024),
                max_args: Some(64),
                max_arg_bytes: Some(64 * 1024),
                max_output_bytes: Some(64 * 1024),
                timeout_s: None,
                max_temp_bytes: None,
                metadata: None,
            }),
        });

        let response = PbAgentService::run_once(&server, request)
            .await
            .unwrap()
            .into_inner();
        let failure = match response.result.expect("a result") {
            super::pb::run_response::Result::Failure(f) => f,
            super::pb::run_response::Result::Success(s) => {
                panic!("unexpected success: {}", s.output)
            }
        };
        assert!(
            failure.message.contains("Roto compilation failed"),
            "expected compile failure message, got: {}",
            failure.message
        );
        assert!(
            failure
                .diagnostics
                .iter()
                .any(|d| d.code == "roto-compile-error"),
            "expected structured diagnostic with code roto-compile-error"
        );
        assert_eq!(failure.kind, pb::ErrorKind::Compile as i32);
    }

    #[test]
    fn map_run_error_preserves_error_kind_and_category() {
        let err = crate::error::Error::resource_exhausted("output too large").with_diagnostic(
            crate::diagnostic::Diagnostic::new(
                crate::diagnostic::Category::ResourceExhausted,
                crate::diagnostic::Severity::Error,
                "output too large",
            )
            .with_code("output-too-large"),
        );

        let response = super::map_run_error(err);
        let failure = match response.result.expect("a result") {
            super::pb::run_response::Result::Failure(f) => f,
            super::pb::run_response::Result::Success(s) => {
                panic!("expected failure, got output: {}", s.output)
            }
        };

        assert_eq!(failure.kind, pb::ErrorKind::ResourceExhausted as i32);
        assert_eq!(
            failure.diagnostics[0].category,
            pb::diagnostic::Category::ResourceExhausted as i32
        );
    }

    #[test]
    fn map_run_error_preserves_cancelled_kind() {
        let err = crate::error::Error::cancelled("timeout exceeded")
            .with_primary_diagnostic("timeout-exceeded");

        let response = super::map_run_error(err);
        let failure = match response.result.expect("a result") {
            super::pb::run_response::Result::Failure(f) => f,
            super::pb::run_response::Result::Success(s) => {
                panic!("expected failure, got output: {}", s.output)
            }
        };

        assert_eq!(failure.kind, pb::ErrorKind::Cancelled as i32);
        assert_eq!(failure.diagnostics[0].code, "timeout-exceeded");
        assert_eq!(
            failure.diagnostics[0].category,
            pb::diagnostic::Category::Cancelled as i32
        );
    }

    #[test]
    fn rpc_cancellation_guard_cancels_when_dropped() {
        let cancellation = RunCancellation::none();
        {
            let _guard = super::RpcCancellationGuard::new(cancellation.clone());
        }
        assert!(cancellation.is_cancelled());
    }

    #[test]
    fn rpc_cancellation_guard_disarm_preserves_token() {
        let cancellation = RunCancellation::none();
        {
            let mut guard = super::RpcCancellationGuard::new(cancellation.clone());
            guard.disarm();
        }
        assert!(!cancellation.is_cancelled());
    }

    #[tokio::test]
    async fn panic_path_returns_run_failure() {
        let panic_module = crate::module::FunctionModule::new(
            crate::module::ModuleSpec {
                name: String::from("panic"),
                doc: String::from("test"),
                capabilities: vec![],
                parameters: vec![],
                functions: vec![],
            },
            |_runtime: &mut roto::Runtime<roto::Ctx<crate::context::ExecutionContext>>, _params| {
                panic!("intentional panic in registration")
            },
            std::sync::Arc::new(|| Ok(())),
        );
        let panic_engine = Arc::new(crate::engine::Engine::new(vec![Arc::new(panic_module)]));
        let panic_service = Arc::new(AgentService::new(
            panic_engine,
            AgentMetadata {
                agent_name: "panic-agent".into(),
                modules: vec![],
                sdk_version: "0.1.0".to_string(),
                declared_limits: std::collections::HashMap::new(),
                metadata: std::collections::HashMap::new(),
            },
        ));
        let panic_server = GrpcServer::new(panic_service);

        let request = tonic::Request::new(RunRequest {
            script: r#"fn main() -> String { "ok" }"#.to_string(),
            args: vec![],
            parameters: std::collections::HashMap::new(),
            limits: None,
        });

        let response = PbAgentService::run_once(&panic_server, request)
            .await
            .unwrap()
            .into_inner();
        let failure = match response.result.expect("a result") {
            super::pb::run_response::Result::Failure(f) => f,
            super::pb::run_response::Result::Success(s) => {
                panic!("expected failure for panic path, got output: {}", s.output)
            }
        };
        assert!(
            failure.message.contains("panic"),
            "expected panic in failure message, got: {}",
            failure.message
        );
        assert_eq!(failure.kind, pb::ErrorKind::Panic as i32);
    }
}
