use std::collections::{BTreeMap, HashMap};
use std::path::{Path, PathBuf};
use std::sync::Arc;

use async_trait::async_trait;

use serde_json::{json, Value};

use crate::client_gen::cli::CliGenerator;
use crate::client_gen::generator::{
    artifact_map, write_artifacts, ClientGenerator, GeneratedArtifact, GeneratorConfig,
};
use crate::client_gen::python::PythonGenerator;
use crate::client_gen::typescript::TypeScriptGenerator;
use crate::compression::engine::Tool;
use crate::compression::CompressionLevel;
use crate::ffi::{normalize_sdk_servers, FfiSdkServerConfig, FfiSdkServersConfig};
use crate::proxy::{RunningToolProxy, ToolProxyServer};
use crate::server::{BackendAuthMode, BackendServerConfig};
use crate::server::{CompressedServer, CompressedServerConfig, ProxyTransformMode};
use crate::Error;

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CompressorMode {
    CompressedTools,
    Cli,
    JustBash,
}

impl From<CompressorMode> for ProxyTransformMode {
    fn from(value: CompressorMode) -> Self {
        match value {
            CompressorMode::CompressedTools => Self::CompressedTools,
            CompressorMode::Cli => Self::Cli,
            CompressorMode::JustBash => Self::JustBash,
        }
    }
}

type HeaderProvider = Arc<dyn Fn() -> Result<BTreeMap<String, String>, Error> + Send + Sync>;

#[derive(Clone)]
pub struct ServerConfig {
    inner: FfiSdkServerConfig,
    auth_provider: Option<HeaderProvider>,
    oauth_app_name: Option<String>,
}

impl ServerConfig {
    pub fn command(command: impl Into<String>) -> Self {
        Self {
            inner: FfiSdkServerConfig::Structured {
                command: Some(command.into()),
                url: None,
                args: Vec::new(),
                headers: BTreeMap::new(),
                oauth_app_name: None,
            },
            auth_provider: None,
            oauth_app_name: None,
        }
    }

    pub fn url(url: impl Into<String>) -> Self {
        Self {
            inner: FfiSdkServerConfig::Structured {
                command: None,
                url: Some(url.into()),
                args: Vec::new(),
                headers: BTreeMap::new(),
                oauth_app_name: None,
            },
            auth_provider: None,
            oauth_app_name: None,
        }
    }

    pub fn arg(mut self, arg: impl Into<String>) -> Self {
        if let FfiSdkServerConfig::Structured { args, .. } = &mut self.inner {
            args.push(arg.into());
        }
        self
    }

    pub fn args(mut self, args: impl IntoIterator<Item = impl Into<String>>) -> Self {
        if let FfiSdkServerConfig::Structured { args: stored, .. } = &mut self.inner {
            stored.extend(args.into_iter().map(Into::into));
        }
        self
    }

    pub fn header(mut self, name: impl Into<String>, value: impl Into<String>) -> Self {
        if let FfiSdkServerConfig::Structured { headers, .. } = &mut self.inner {
            headers.insert(name.into(), value.into());
        }
        self
    }

    pub fn auth_provider(
        mut self,
        provider: impl Fn() -> Result<BTreeMap<String, String>, Error> + Send + Sync + 'static,
    ) -> Self {
        self.auth_provider = Some(Arc::new(provider));
        self
    }

    pub fn oauth_app_name(mut self, app_name: impl Into<String>) -> Self {
        self.oauth_app_name = Some(app_name.into());
        self
    }

    fn materialize(mut self) -> (FfiSdkServerConfig, Option<HeaderProvider>) {
        if let (FfiSdkServerConfig::Structured { oauth_app_name, .. }, Some(app_name)) =
            (&mut self.inner, self.oauth_app_name.take())
        {
            *oauth_app_name = Some(app_name);
        }
        (self.inner, self.auth_provider.take())
    }
}

#[derive(Clone)]
pub struct CompressorClientBuilder {
    servers: BTreeMap<String, ServerConfig>,
    compression_level: CompressionLevel,
    server_name: Option<String>,
    include_tools: Vec<String>,
    exclude_tools: Vec<String>,
    toonify: bool,
    mode: CompressorMode,
}

impl Default for CompressorClientBuilder {
    fn default() -> Self {
        Self {
            servers: BTreeMap::new(),
            compression_level: CompressionLevel::Max,
            server_name: None,
            include_tools: Vec::new(),
            exclude_tools: Vec::new(),
            toonify: false,
            mode: CompressorMode::CompressedTools,
        }
    }
}

impl CompressorClientBuilder {
    pub fn server(mut self, name: impl Into<String>, config: ServerConfig) -> Self {
        self.servers.insert(name.into(), config);
        self
    }

    pub fn compression_level(mut self, level: CompressionLevel) -> Self {
        self.compression_level = level;
        self
    }

    pub fn server_name(mut self, server_name: impl Into<String>) -> Self {
        self.server_name = Some(server_name.into());
        self
    }

    pub fn include_tools(mut self, tools: impl IntoIterator<Item = impl Into<String>>) -> Self {
        self.include_tools = tools.into_iter().map(Into::into).collect();
        self
    }

    pub fn exclude_tools(mut self, tools: impl IntoIterator<Item = impl Into<String>>) -> Self {
        self.exclude_tools = tools.into_iter().map(Into::into).collect();
        self
    }

    pub fn toonify(mut self, enabled: bool) -> Self {
        self.toonify = enabled;
        self
    }

    pub fn mode(mut self, mode: CompressorMode) -> Self {
        self.mode = mode;
        self
    }

    pub fn build(self) -> CompressorClient {
        CompressorClient { builder: self }
    }
}

#[derive(Clone)]
pub struct CompressorClient {
    builder: CompressorClientBuilder,
}

impl CompressorClient {
    pub fn builder() -> CompressorClientBuilder {
        CompressorClientBuilder::default()
    }

    pub async fn connect(&self) -> Result<CompressorProxy, Error> {
        let materialized = self
            .builder
            .servers
            .clone()
            .into_iter()
            .map(|(name, config)| {
                let (config, provider) = config.materialize();
                (name, config, provider)
            })
            .collect::<Vec<_>>();
        let providers = materialized
            .iter()
            .filter_map(|(name, _, provider)| {
                provider.clone().map(|provider| (name.clone(), provider))
            })
            .collect::<BTreeMap<_, _>>();
        let ffi_configs = materialized
            .into_iter()
            .map(|(name, config, _)| (name, config))
            .collect::<Vec<_>>();
        let backends = normalize_sdk_servers(FfiSdkServersConfig::from_iter(ffi_configs))?;
        let backends = backends
            .into_iter()
            .map(|backend| {
                let name = backend.name.clone();
                let mut backend = BackendServerConfig::from(backend);
                if let Some(provider) = providers.get(&name) {
                    backend = backend
                        .with_header_provider(Arc::clone(provider))
                        .with_auth_mode(BackendAuthMode::ExplicitHeaders);
                }
                backend
            })
            .collect::<Vec<_>>();
        let server = CompressedServer::connect_multi_stdio(
            CompressedServerConfig {
                level: self.builder.compression_level.clone(),
                server_name: self.builder.server_name.clone(),
                include_tools: self.builder.include_tools.clone(),
                exclude_tools: self.builder.exclude_tools.clone(),
                toonify: self.builder.toonify,
                transform_mode: self.builder.mode.into(),
                ..CompressedServerConfig::default()
            },
            backends,
        )
        .await?;
        CompressorProxy::start(server).await
    }
}

pub struct CompressorProxy {
    default_server: Option<String>,
    frontend_tools: Vec<Tool>,
    backend_tools: Vec<Tool>,
    backend_tools_by_server: Vec<(String, Tool)>,
    just_bash_providers: Vec<crate::server::JustBashProviderSpec>,
    proxy: RunningToolProxy,
}

impl CompressorProxy {
    async fn start(server: CompressedServer) -> Result<Self, Error> {
        let default_server = server.default_server_name().map(str::to_string);
        let frontend_tools = server.list_frontend_tools().await?;
        let backend_tools = server.backend_tools();
        let backend_tools_by_server = server.backend_tools_by_server();
        let just_bash_providers = server.just_bash_provider_specs();
        let proxy = ToolProxyServer::start(server).await?;
        Ok(Self {
            default_server,
            frontend_tools,
            backend_tools,
            backend_tools_by_server,
            just_bash_providers,
            proxy,
        })
    }

    pub fn bridge_url(&self) -> &str {
        self.proxy.bridge_url()
    }

    pub fn token(&self) -> &str {
        self.proxy.token_value()
    }

    pub fn tools(&self) -> &[Tool] {
        &self.frontend_tools
    }

    pub fn backend_tools(&self) -> &[Tool] {
        &self.backend_tools
    }

    pub fn just_bash_providers(&self) -> &[crate::server::JustBashProviderSpec] {
        &self.just_bash_providers
    }

    pub fn schema(&self, tool_name: &str) -> Result<Value, Error> {
        self.schema_on(self.default_server.as_deref(), tool_name)
    }

    pub fn schema_on(&self, server: Option<&str>, tool_name: &str) -> Result<Value, Error> {
        let matches = self
            .backend_tools_by_server
            .iter()
            .filter(|(server_name, tool)| {
                tool.name == tool_name && server.map(|server| server == server_name).unwrap_or(true)
            })
            .collect::<Vec<_>>();
        match matches.as_slice() {
            [(_, tool)] => Ok(tool.input_schema.clone()),
            [] => Err(Error::ToolNotFound(tool_name.to_string())),
            _ => Err(Error::Config(
                "Multiple backend tools matched; specify a server".to_string(),
            )),
        }
    }

    pub async fn invoke(&self, tool_name: &str, input: Value) -> Result<String, Error> {
        self.invoke_on(self.default_server.as_deref(), tool_name, input)
            .await
    }

    pub async fn invoke_on(
        &self,
        server: Option<&str>,
        tool_name: &str,
        input: Value,
    ) -> Result<String, Error> {
        let wrapper = self.invoke_wrapper(server)?;
        self.invoke_wrapper_tool(
            &wrapper,
            json!({
                "tool_name": tool_name,
                "tool_input": input,
            }),
        )
        .await
    }

    async fn invoke_wrapper_tool(&self, wrapper: &str, input: Value) -> Result<String, Error> {
        let exec_url = self.proxy.exec_url();
        ensure_loopback_proxy_url(&exec_url)?;
        let client = reqwest::Client::new();
        let response = client
            // codeql[rust/cleartext-transmission]
            // This session token is sent only to the local loopback proxy after
            // ensure_loopback_proxy_url rejects non-loopback and non-http URLs.
            .post(exec_url)
            .header("Authorization", format!("Bearer {}", self.token()))
            .json(&json!({
                "tool": wrapper,
                "input": input
            }))
            .send()
            .await
            .map_err(|error| Error::Config(format!("proxy request failed: {error}")))?;
        let status = response.status();
        let text = response
            .text()
            .await
            .map_err(|error| Error::Config(format!("proxy response failed: {error}")))?;
        if status.is_success() {
            Ok(text)
        } else {
            Err(Error::Config(format!(
                "proxy request failed with {status}: {text}"
            )))
        }
    }

    pub fn executable_tools(&self) -> BTreeMap<String, Box<dyn ExecutableTool + '_>> {
        self.frontend_tools
            .iter()
            .map(|tool| {
                (
                    tool.name.clone(),
                    Box::new(ProxyExecutableTool {
                        proxy: self,
                        tool: tool.clone(),
                    }) as Box<dyn ExecutableTool>,
                )
            })
            .collect()
    }

    pub fn write_cli_client(
        &self,
        output_dir: impl AsRef<Path>,
        name: Option<&str>,
    ) -> Result<GeneratedClient, Error> {
        self.write_client(GeneratedClientKind::Cli, output_dir, name)
    }

    pub fn generate_code_client(
        &self,
        language: CodeLanguage,
        output_dir: impl AsRef<Path>,
        name: Option<&str>,
    ) -> Result<GeneratedClient, Error> {
        let kind = match language {
            CodeLanguage::Python => GeneratedClientKind::Python,
            CodeLanguage::TypeScript => GeneratedClientKind::TypeScript,
        };
        self.generate_client(kind, output_dir, name)
    }

    pub fn write_code_client(
        &self,
        language: CodeLanguage,
        output_dir: impl AsRef<Path>,
        name: Option<&str>,
    ) -> Result<GeneratedClient, Error> {
        let generated = self.generate_code_client(language, output_dir, name)?;
        generated.write_to_disk()
    }

    pub fn generate_client(
        &self,
        kind: GeneratedClientKind,
        output_dir: impl AsRef<Path>,
        name: Option<&str>,
    ) -> Result<GeneratedClient, Error> {
        let generator_config = self.generator_config(output_dir, name);
        let artifacts = match kind {
            GeneratedClientKind::Cli => CliGenerator.render(&generator_config),
            GeneratedClientKind::Python => PythonGenerator.render(&generator_config),
            GeneratedClientKind::TypeScript => TypeScriptGenerator.render(&generator_config),
        }?;
        let file_contents = artifact_map(&artifacts);
        let files = artifacts
            .iter()
            .map(|artifact| generator_config.output_dir.join(&artifact.file_name))
            .collect();
        let environment = kind.environment(&generator_config);
        Ok(GeneratedClient {
            kind,
            output_dir: generator_config.output_dir,
            files,
            file_contents,
            artifacts,
            environment,
        })
    }

    pub fn write_client(
        &self,
        kind: GeneratedClientKind,
        output_dir: impl AsRef<Path>,
        name: Option<&str>,
    ) -> Result<GeneratedClient, Error> {
        let generated = self.generate_client(kind, output_dir, name)?;
        generated.write_to_disk()
    }

    fn generator_config(
        &self,
        output_dir: impl AsRef<Path>,
        name: Option<&str>,
    ) -> GeneratorConfig {
        GeneratorConfig {
            cli_name: name
                .or(self.default_server.as_deref())
                .unwrap_or("mcp")
                .to_string(),
            bridge_url: self.bridge_url().to_string(),
            token: self.token().to_string(),
            tools: self.backend_tools.clone(),
            session_pid: 0,
            output_dir: output_dir.as_ref().to_path_buf(),
            extra_cli_bridges: Vec::new(),
        }
    }

    fn invoke_wrapper(&self, server: Option<&str>) -> Result<String, Error> {
        let suffix = "_invoke_tool";
        let matches = self
            .frontend_tools
            .iter()
            .filter(|tool| tool.name.ends_with(suffix))
            .filter(|tool| {
                server
                    .map(|name| tool.name == format!("{name}{suffix}"))
                    .unwrap_or(true)
            })
            .map(|tool| tool.name.clone())
            .collect::<Vec<_>>();
        match matches.as_slice() {
            [name] => Ok(name.clone()),
            [] => Err(Error::Config(format!(
                "No compressed invoke wrapper found for server {}",
                server.unwrap_or("<default>")
            ))),
            _ => Err(Error::Config(
                "Multiple compressed invoke wrappers found; specify a server".to_string(),
            )),
        }
    }
}

fn ensure_loopback_proxy_url(url: &str) -> Result<(), Error> {
    let parsed = reqwest::Url::parse(url)
        .map_err(|error| Error::Config(format!("invalid proxy URL {url:?}: {error}")))?;
    if parsed.scheme() != "http" {
        return Err(Error::Config(format!(
            "refusing to send proxy session token to non-http URL: {url}"
        )));
    }
    let Some(host) = parsed.host_str() else {
        return Err(Error::Config(format!("proxy URL has no host: {url}")));
    };
    let is_loopback_ip = host
        .parse::<std::net::IpAddr>()
        .map(|addr| addr.is_loopback())
        .unwrap_or(false);
    if host != "localhost" && !is_loopback_ip {
        return Err(Error::Config(format!(
            "refusing to send proxy session token to non-loopback URL: {url}"
        )));
    }
    Ok(())
}

#[async_trait]
pub trait ExecutableTool: Send + Sync {
    fn name(&self) -> &str;
    fn description(&self) -> Option<&str>;
    fn input_schema(&self) -> &Value;
    async fn execute(&self, input: Value) -> Result<String, Error>;
}

struct ProxyExecutableTool<'a> {
    proxy: &'a CompressorProxy,
    tool: Tool,
}

#[async_trait]
impl ExecutableTool for ProxyExecutableTool<'_> {
    fn name(&self) -> &str {
        &self.tool.name
    }

    fn description(&self) -> Option<&str> {
        self.tool.description.as_deref()
    }

    fn input_schema(&self) -> &Value {
        &self.tool.input_schema
    }

    async fn execute(&self, input: Value) -> Result<String, Error> {
        self.proxy.invoke_wrapper_tool(&self.tool.name, input).await
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CodeLanguage {
    Python,
    TypeScript,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct GeneratedClient {
    pub kind: GeneratedClientKind,
    pub output_dir: PathBuf,
    /// Paths where artifacts would be written, or were written after `write_to_disk`.
    pub files: Vec<PathBuf>,
    /// Generated artifact contents keyed by file name.
    pub file_contents: BTreeMap<String, String>,
    pub environment: HashMap<String, String>,
    artifacts: Vec<GeneratedArtifact>,
}

impl GeneratedClient {
    pub fn write_to_disk(mut self) -> Result<Self, Error> {
        self.files = write_artifacts(&self.artifacts, &self.output_dir)?;
        Ok(self)
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum GeneratedClientKind {
    Cli,
    Python,
    TypeScript,
}

impl GeneratedClientKind {
    fn environment(self, config: &GeneratorConfig) -> HashMap<String, String> {
        match self {
            GeneratedClientKind::Python => HashMap::from([(
                "PYTHONPATH".to_string(),
                config.output_dir.to_string_lossy().to_string(),
            )]),
            GeneratedClientKind::Cli | GeneratedClientKind::TypeScript => HashMap::new(),
        }
    }
}

#[cfg(test)]
mod tests {
    use serde_json::json;

    use super::*;

    fn fixture_path(name: &str) -> String {
        format!("{}/tests/fixtures/{name}", env!("CARGO_MANIFEST_DIR"))
    }

    fn python_command() -> String {
        std::env::var("PYTHON").unwrap_or_else(|_| "python3".to_string())
    }

    #[test]
    fn server_config_oauth_app_name_is_preserved_for_transport_layer() {
        let config = ServerConfig::url("https://example.test/mcp")
            .oauth_app_name("Rovo Dev")
            .materialize()
            .0;

        match config {
            FfiSdkServerConfig::Structured { oauth_app_name, .. } => {
                assert_eq!(oauth_app_name.as_deref(), Some("Rovo Dev"));
            }
            FfiSdkServerConfig::CommandOrUrl(_) => panic!("expected structured config"),
        }
    }

    #[test]
    fn loopback_proxy_url_guard_allows_only_local_http_urls() {
        assert!(ensure_loopback_proxy_url("http://127.0.0.1:1234/exec").is_ok());
        assert!(ensure_loopback_proxy_url("http://localhost:1234/exec").is_ok());
        assert!(ensure_loopback_proxy_url("https://127.0.0.1:1234/exec").is_err());
        assert!(ensure_loopback_proxy_url("http://example.com:1234/exec").is_err());
    }

    #[test]
    fn server_config_auth_provider_is_preserved_for_transport_layer() {
        let (config, provider) = ServerConfig::url("https://example.test/mcp")
            .header("X-Static", "yes")
            .auth_provider(|| {
                Ok(BTreeMap::from([(
                    "Authorization".to_string(),
                    "Bearer dynamic".to_string(),
                )]))
            })
            .materialize();

        let backends = normalize_sdk_servers(FfiSdkServersConfig::from_iter([(
            "remote".to_string(),
            config,
        )]))
        .unwrap();

        assert_eq!(backends[0].command_or_url, "https://example.test/mcp");
        assert_eq!(
            backends[0].args,
            ["-H", "X-Static=yes", "--auth", "explicit-headers"]
        );
        assert!(provider.is_some());
    }

    #[tokio::test]
    async fn compressor_client_invokes_single_server_without_compressor_subprocess() {
        let client = CompressorClient::builder()
            .server(
                "alpha",
                ServerConfig::command(python_command()).arg(fixture_path("alpha_server.py")),
            )
            .compression_level(CompressionLevel::Max)
            .build();
        let proxy = client.connect().await.unwrap();
        assert!(proxy
            .tools()
            .iter()
            .any(|tool| tool.name == "alpha_invoke_tool"));
        let result = proxy
            .invoke("echo", json!({ "message": "rust-sdk" }))
            .await
            .unwrap();
        assert_eq!(result, "alpha:rust-sdk");

        let executable = proxy.executable_tools();
        let invoke = executable.get("alpha_invoke_tool").unwrap();
        let executable_result = invoke
            .execute(json!({
                "tool_name": "echo",
                "tool_input": { "message": "executable-rust" }
            }))
            .await
            .unwrap();
        assert_eq!(executable_result, "alpha:executable-rust");
    }

    #[tokio::test]
    async fn compressor_client_routes_multiple_servers() {
        let client = CompressorClient::builder()
            .server(
                "alpha",
                ServerConfig::command(python_command()).arg(fixture_path("alpha_server.py")),
            )
            .server(
                "beta",
                ServerConfig::command(python_command()).arg(fixture_path("beta_server.py")),
            )
            .compression_level(CompressionLevel::Max)
            .build();
        let proxy = client.connect().await.unwrap();
        let alpha = proxy
            .invoke_on(Some("alpha"), "add", json!({ "a": 2, "b": 3 }))
            .await
            .unwrap();
        let beta = proxy
            .invoke_on(Some("beta"), "multiply", json!({ "a": 4, "b": 5 }))
            .await
            .unwrap();
        assert_eq!(alpha, "5");
        assert_eq!(beta, "20");
    }

    #[tokio::test]
    async fn compressor_client_writes_generated_clients() {
        let client = CompressorClient::builder()
            .server(
                "alpha",
                ServerConfig::command(python_command()).arg(fixture_path("alpha_server.py")),
            )
            .compression_level(CompressionLevel::Max)
            .build();
        let proxy = client.connect().await.unwrap();
        let tempdir = tempfile::tempdir().unwrap();
        let generated = proxy
            .write_code_client(CodeLanguage::Python, tempdir.path(), Some("alpha"))
            .unwrap();
        assert_eq!(generated.kind, GeneratedClientKind::Python);
        assert!(generated
            .files
            .iter()
            .any(|path| path.ends_with("alpha.py")));
        assert_eq!(
            generated.environment.get("PYTHONPATH"),
            Some(&tempdir.path().to_string_lossy().to_string())
        );

        let cli = proxy
            .write_cli_client(tempdir.path(), Some("alpha"))
            .unwrap();
        assert_eq!(cli.kind, GeneratedClientKind::Cli);
    }

    #[tokio::test]
    async fn compressor_client_exposes_cli_and_just_bash_modes() {
        let cli = CompressorClient::builder()
            .server(
                "alpha",
                ServerConfig::command(python_command()).arg(fixture_path("alpha_server.py")),
            )
            .mode(CompressorMode::Cli)
            .build()
            .connect()
            .await
            .unwrap();
        assert!(cli.tools().iter().any(|tool| tool.name == "alpha_help"));

        let bash = CompressorClient::builder()
            .server(
                "alpha",
                ServerConfig::command(python_command()).arg(fixture_path("alpha_server.py")),
            )
            .mode(CompressorMode::JustBash)
            .build()
            .connect()
            .await
            .unwrap();
        assert!(bash.tools().iter().any(|tool| tool.name == "bash_tool"));
        assert!(bash.tools().iter().any(|tool| tool.name == "alpha_help"));
        let provider = bash
            .just_bash_providers()
            .iter()
            .find(|provider| provider.provider_name == "alpha")
            .unwrap();
        assert_eq!(provider.help_tool_name, "alpha_help");
        assert!(provider.tools.iter().any(|command| {
            command.command_name == "echo"
                && command.backend_tool_name == "echo"
                && command.invoke_tool_name == "alpha_invoke_tool"
        }));
    }
}
