mod client;
mod discovery;
mod policy;

use std::fs;
use std::io::BufReader;
use std::path::PathBuf;
use std::sync::Arc;
use std::sync::Once;

use anyhow::{anyhow, Context, Result};
use clap::Parser;
use rustls::{pki_types::CertificateDer, pki_types::PrivateKeyDer, ServerConfig};
use rustls_pemfile::{certs, pkcs8_private_keys, rsa_private_keys};
use serde::Deserialize;
use tokio::io;
use tokio::net::{TcpListener, TcpStream};
use tokio::task;
use tokio_rustls::TlsAcceptor;
use tracing::{info, warn};
use tracing_subscriber::EnvFilter;

use crate::discovery::ServiceDiscovery;
use crate::policy::PolicyEnforcer;

/// mTLS Agent (MVP: プレーンTCPプロキシ)
///
/// - 設計どおり、アプリケーションは `app_addr` で平文待受
/// - 本Agentは `mesh_bind_addr` で待受し、受信した接続を `app_addr` にフォワードするだけ
/// - mTLS/TLS 対応は後続で追加する前提のスケルトン
#[derive(Parser, Debug)]
#[command(author, version, about)]
struct Cli {
    /// 設定ファイル (TOML)
    #[arg(long)]
    config: PathBuf,
}

#[derive(Debug, Deserialize)]
struct ServiceConfig {
    name: String,
    app_addr: String,
    mesh_bind_addr: String,
}

#[derive(Debug, Deserialize)]
struct ClusterConfig {
    cluster_id: String,
    environment: Option<String>,
    chainfire_endpoint: Option<String>,
}

#[derive(Debug, Deserialize)]
struct MtlsConfig {
    #[serde(default)]
    mode: Option<String>, // auto/mtls/tls/plain
    #[serde(default)]
    ca_cert_path: Option<String>,
    #[serde(default)]
    cert_path: Option<String>,
    #[serde(default)]
    key_path: Option<String>,
}

#[derive(Debug, Deserialize)]
struct Config {
    service: ServiceConfig,
    #[serde(default)]
    cluster: Option<ClusterConfig>,
    #[serde(default)]
    mtls: Option<MtlsConfig>,
}

#[tokio::main]
async fn main() -> Result<()> {
    tracing_subscriber::fmt()
        .with_env_filter(EnvFilter::from_default_env().add_directive("info".parse()?))
        .init();

    let cli = Cli::parse();
    let cfg = load_config(&cli.config)?;

    let mode = cfg
        .mtls
        .as_ref()
        .and_then(|m| m.mode.as_deref())
        .unwrap_or("plain")
        .to_lowercase();

    info!(
        service = %cfg.service.name,
        mesh_bind = %cfg.service.mesh_bind_addr,
        app_addr = %cfg.service.app_addr,
        mode = %mode,
        "starting mtls-agent"
    );

    // Chainfire統合: サービス発見とポリシー管理
    let (discovery, _policy_enforcer) = if let Some(cluster_cfg) = &cfg.cluster {
        if let Some(endpoint) = &cluster_cfg.chainfire_endpoint {
            let disc = Arc::new(ServiceDiscovery::new(
                endpoint.clone(),
                cluster_cfg.cluster_id.clone(),
            ));
            disc.start_background_refresh().await;

            let enforcer = PolicyEnforcer::new(Arc::clone(&disc), mode.clone());
            enforcer.start_background_refresh().await;

            (Some(disc), Some(enforcer))
        } else {
            (None, None)
        }
    } else {
        (None, None)
    };

    // モード決定: "auto" の場合はChainfireからポリシーを読み取る
    let effective_mode = if mode == "auto" {
        if let Some(disc) = &discovery {
            // デフォルトポリシーを確認（簡易実装）
            // 実際には、自身のサービス名とターゲットサービス名でポリシーを検索
            if let Ok(Some(policy)) = disc.get_mtls_policy(&cfg.service.name, "default").await {
                if policy.mtls_required.unwrap_or(false) {
                    "mtls"
                } else {
                    "tls"
                }
            } else {
                // 環境に応じたデフォルト
                let env = cfg
                    .cluster
                    .as_ref()
                    .and_then(|c| c.environment.as_deref())
                    .unwrap_or("dev");
                if env == "prod" || env == "stg" {
                    "mtls"
                } else {
                    "plain"
                }
            }
        } else {
            "plain"
        }
    } else {
        mode.as_str()
    };

    info!(effective_mode = %effective_mode, "determined mTLS mode");

    match effective_mode {
        "plain" => {
            run_plain_proxy(&cfg.service.mesh_bind_addr, &cfg.service.app_addr).await?;
        }
        "tls" | "mtls" => {
            let tls_cfg = build_server_config(&cfg, effective_mode)?;
            run_tls_proxy(&cfg.service.mesh_bind_addr, &cfg.service.app_addr, tls_cfg).await?;
        }
        other => {
            return Err(anyhow!("unsupported mtls.mode: {}", other));
        }
    }

    Ok(())
}

fn load_config(path: &PathBuf) -> Result<Config> {
    let contents =
        fs::read_to_string(path).with_context(|| format!("failed to read {}", path.display()))?;
    let cfg: Config =
        toml::from_str(&contents).with_context(|| format!("failed to parse {}", path.display()))?;
    Ok(cfg)
}

fn load_certs(path: &str) -> Result<Vec<CertificateDer<'static>>> {
    let file =
        fs::File::open(path).with_context(|| format!("failed to open cert file {}", path))?;
    let mut reader = BufReader::new(file);
    let certs = certs(&mut reader)
        .collect::<std::result::Result<Vec<_>, _>>()
        .map_err(|e| anyhow!("failed to parse certs from {}: {}", path, e))?;
    Ok(certs)
}

fn load_private_key(path: &str) -> Result<PrivateKeyDer<'static>> {
    let file = fs::File::open(path).with_context(|| format!("failed to open key file {}", path))?;
    let mut reader = BufReader::new(file);

    // Try PKCS8 first
    if let Ok(keys) = pkcs8_private_keys(&mut reader).collect::<std::result::Result<Vec<_>, _>>() {
        if let Some(k) = keys.into_iter().next() {
            return Ok(PrivateKeyDer::Pkcs8(k));
        }
    }

    // Fallback to RSA
    let file = fs::File::open(path).with_context(|| format!("failed to open key file {}", path))?;
    let mut reader = BufReader::new(file);
    let keys = rsa_private_keys(&mut reader)
        .collect::<std::result::Result<Vec<_>, _>>()
        .map_err(|e| anyhow!("failed to parse private key from {}: {}", path, e))?;
    let Some(k) = keys.into_iter().next() else {
        return Err(anyhow!("no private keys found in {}", path));
    };
    Ok(PrivateKeyDer::Pkcs1(k))
}

fn ensure_crypto_provider() {
    static INIT: Once = Once::new();
    INIT.call_once(|| {
        let _ = rustls::crypto::aws_lc_rs::default_provider().install_default();
    });
}

fn build_server_config(cfg: &Config, mode: &str) -> Result<ServerConfig> {
    ensure_crypto_provider();

    let mtls = cfg
        .mtls
        .as_ref()
        .ok_or_else(|| anyhow!("mtls section is required for mode {}", mode))?;

    let cert_path = mtls
        .cert_path
        .as_deref()
        .ok_or_else(|| anyhow!("mtls.cert_path is required"))?;
    let key_path = mtls
        .key_path
        .as_deref()
        .ok_or_else(|| anyhow!("mtls.key_path is required"))?;

    let certs = load_certs(cert_path)?;
    let key = load_private_key(key_path)?;

    // ベースは「クライアント認証なし」のサーバ設定
    let builder = ServerConfig::builder();

    // mTLS の場合はクライアント証明書を要求する verifer を使う
    if mode == "mtls" {
        let ca_path = mtls
            .ca_cert_path
            .as_deref()
            .ok_or_else(|| anyhow!("mtls.ca_cert_path is required for mtls mode"))?;
        let client_certs = load_certs(ca_path)?;
        let mut roots = rustls::RootCertStore::empty();
        for c in client_certs {
            roots
                .add(c)
                .map_err(|e| anyhow!("adding CA failed: {:?}", e))?;
        }
        let verifier =
            rustls::server::WebPkiClientVerifier::builder(std::sync::Arc::new(roots)).build()?;
        let cfg = builder
            .with_client_cert_verifier(verifier)
            .with_single_cert(certs, key)
            .map_err(|e| anyhow!("failed to build mtls server config: {}", e))?;
        Ok(cfg)
    } else {
        let cfg = builder
            .with_no_client_auth()
            .with_single_cert(certs, key)
            .map_err(|e| anyhow!("failed to build tls server config: {}", e))?;
        Ok(cfg)
    }
}

async fn run_plain_proxy(listen_addr: &str, app_addr: &str) -> Result<()> {
    let listener = TcpListener::bind(listen_addr).await?;
    info!(
        "listening on {} and forwarding to {}",
        listen_addr, app_addr
    );

    loop {
        let (inbound, peer) = listener.accept().await?;
        let app_addr = app_addr.to_string();

        info!(remote = %peer, "accepted connection");

        task::spawn(async move {
            if let Err(e) = handle_connection(inbound, &app_addr).await {
                warn!(error = %e, "connection handling failed");
            }
        });
    }
}

async fn run_tls_proxy(
    listen_addr: &str,
    app_addr: &str,
    server_config: ServerConfig,
) -> Result<()> {
    let listener = TcpListener::bind(listen_addr).await?;
    let acceptor = TlsAcceptor::from(std::sync::Arc::new(server_config));

    info!(
        "listening (TLS/mTLS) on {} and forwarding to {}",
        listen_addr, app_addr
    );

    loop {
        let (inbound, peer) = listener.accept().await?;
        let app_addr = app_addr.to_string();
        let acceptor = acceptor.clone();

        info!(remote = %peer, "accepted TLS connection");

        task::spawn(async move {
            match acceptor.accept(inbound).await {
                Ok(tls_stream) => {
                    if let Err(e) = handle_tls_connection(tls_stream, &app_addr).await {
                        warn!(error = %e, "TLS connection handling failed");
                    }
                }
                Err(e) => {
                    warn!(error = %e, "TLS handshake failed");
                }
            }
        });
    }
}

async fn handle_tls_connection<S>(inbound: S, app_addr: &str) -> Result<()>
where
    S: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin + Send + 'static,
{
    let mut outbound = TcpStream::connect(app_addr).await?;
    let (mut ri, mut wi) = tokio::io::split(inbound);
    let (mut ro, mut wo) = outbound.split();

    let client_to_app = io::copy(&mut ri, &mut wo);
    let app_to_client = io::copy(&mut ro, &mut wi);

    tokio::try_join!(client_to_app, app_to_client)?;
    Ok(())
}

async fn handle_connection(mut inbound: TcpStream, app_addr: &str) -> Result<()> {
    let mut outbound = TcpStream::connect(app_addr).await?;
    let (mut ri, mut wi) = inbound.split();
    let (mut ro, mut wo) = outbound.split();

    let client_to_app = io::copy(&mut ri, &mut wo);
    let app_to_client = io::copy(&mut ro, &mut wi);

    tokio::try_join!(client_to_app, app_to_client)?;
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use rustls::pki_types::ServerName;
    use rustls::{ClientConfig, RootCertStore};
    use std::path::{Path, PathBuf};
    use std::process::Command;
    use std::sync::OnceLock;
    use tokio::io::{AsyncReadExt, AsyncWriteExt};
    use tokio::time::{sleep, Duration};
    use tokio_rustls::TlsConnector;

    fn run_openssl(args: &[&str]) {
        let status = Command::new("openssl")
            .args(args)
            .status()
            .unwrap_or_else(|error| panic!("failed to spawn openssl {:?}: {}", args, error));
        assert!(
            status.success(),
            "openssl {:?} failed with status {}",
            args,
            status
        );
    }

    fn ensure_test_certs() -> &'static Path {
        static CERT_DIR: OnceLock<PathBuf> = OnceLock::new();

        CERT_DIR.get_or_init(|| {
            let dir =
                std::env::temp_dir().join(format!("mtls-agent-test-certs-{}", std::process::id()));
            std::fs::create_dir_all(&dir).unwrap();

            let ca_key = dir.join("ca.key");
            let ca_pem = dir.join("ca.pem");
            let ca_srl = dir.join("ca.srl");
            let server_key = dir.join("server.key");
            let server_csr = dir.join("server.csr");
            let server_ext = dir.join("server.ext");
            let server_pem = dir.join("server.pem");
            let client_key = dir.join("client.key");
            let client_csr = dir.join("client.csr");
            let client_ext = dir.join("client.ext");
            let client_pem = dir.join("client.pem");

            if !ca_pem.exists() {
                run_openssl(&["genrsa", "-out", ca_key.to_string_lossy().as_ref(), "2048"]);
                run_openssl(&[
                    "req",
                    "-x509",
                    "-new",
                    "-key",
                    ca_key.to_string_lossy().as_ref(),
                    "-sha256",
                    "-days",
                    "3650",
                    "-subj",
                    "/CN=UltraCloud Test CA",
                    "-out",
                    ca_pem.to_string_lossy().as_ref(),
                ]);

                std::fs::write(
                    &server_ext,
                    "subjectAltName=DNS:localhost\nextendedKeyUsage=serverAuth\n",
                )
                .unwrap();
                run_openssl(&[
                    "genrsa",
                    "-out",
                    server_key.to_string_lossy().as_ref(),
                    "2048",
                ]);
                run_openssl(&[
                    "req",
                    "-new",
                    "-key",
                    server_key.to_string_lossy().as_ref(),
                    "-subj",
                    "/CN=localhost",
                    "-out",
                    server_csr.to_string_lossy().as_ref(),
                ]);
                run_openssl(&[
                    "x509",
                    "-req",
                    "-in",
                    server_csr.to_string_lossy().as_ref(),
                    "-CA",
                    ca_pem.to_string_lossy().as_ref(),
                    "-CAkey",
                    ca_key.to_string_lossy().as_ref(),
                    "-CAcreateserial",
                    "-CAserial",
                    ca_srl.to_string_lossy().as_ref(),
                    "-out",
                    server_pem.to_string_lossy().as_ref(),
                    "-days",
                    "3650",
                    "-sha256",
                    "-extfile",
                    server_ext.to_string_lossy().as_ref(),
                ]);

                std::fs::write(&client_ext, "extendedKeyUsage=clientAuth\n").unwrap();
                run_openssl(&[
                    "genrsa",
                    "-out",
                    client_key.to_string_lossy().as_ref(),
                    "2048",
                ]);
                run_openssl(&[
                    "req",
                    "-new",
                    "-key",
                    client_key.to_string_lossy().as_ref(),
                    "-subj",
                    "/CN=ultracloud-test-client",
                    "-out",
                    client_csr.to_string_lossy().as_ref(),
                ]);
                run_openssl(&[
                    "x509",
                    "-req",
                    "-in",
                    client_csr.to_string_lossy().as_ref(),
                    "-CA",
                    ca_pem.to_string_lossy().as_ref(),
                    "-CAkey",
                    ca_key.to_string_lossy().as_ref(),
                    "-CAserial",
                    ca_srl.to_string_lossy().as_ref(),
                    "-out",
                    client_pem.to_string_lossy().as_ref(),
                    "-days",
                    "3650",
                    "-sha256",
                    "-extfile",
                    client_ext.to_string_lossy().as_ref(),
                ]);
            }

            dir
        })
    }

    fn test_cert_path(name: &str) -> String {
        ensure_test_certs().join(name).display().to_string()
    }

    fn unused_loopback_addr() -> String {
        let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
        let addr = listener.local_addr().unwrap();
        drop(listener);
        addr.to_string()
    }

    async fn wait_for_listener(addr: &str) {
        for _ in 0..50 {
            if TcpStream::connect(addr).await.is_ok() {
                return;
            }
            sleep(Duration::from_millis(50)).await;
        }
        panic!("timed out waiting for listener at {}", addr);
    }

    async fn spawn_echo_server() -> String {
        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
        let addr = listener.local_addr().unwrap().to_string();
        tokio::spawn(async move {
            loop {
                let (mut socket, _) = listener.accept().await.unwrap();
                tokio::spawn(async move {
                    let mut buf = vec![0u8; 4096];
                    let n = socket.read(&mut buf).await.unwrap();
                    if n > 0 {
                        socket.write_all(&buf[..n]).await.unwrap();
                    }
                });
            }
        });
        addr
    }

    fn tls_config(mode: &str) -> Config {
        Config {
            service: ServiceConfig {
                name: "mtls-agent-test".to_string(),
                app_addr: "127.0.0.1:1".to_string(),
                mesh_bind_addr: "127.0.0.1:2".to_string(),
            },
            cluster: None,
            mtls: Some(MtlsConfig {
                mode: Some(mode.to_string()),
                ca_cert_path: Some(test_cert_path("ca.pem")),
                cert_path: Some(test_cert_path("server.pem")),
                key_path: Some(test_cert_path("server.key")),
            }),
        }
    }

    fn root_store() -> RootCertStore {
        ensure_crypto_provider();

        let mut roots = RootCertStore::empty();
        for cert in load_certs(&test_cert_path("ca.pem")).unwrap() {
            roots.add(cert).unwrap();
        }
        roots
    }

    #[tokio::test]
    async fn plain_proxy_forwards_traffic() {
        let backend_addr = spawn_echo_server().await;
        let proxy_addr = unused_loopback_addr();

        let proxy = tokio::spawn({
            let proxy_addr = proxy_addr.clone();
            let backend_addr = backend_addr.clone();
            async move {
                run_plain_proxy(&proxy_addr, &backend_addr).await.unwrap();
            }
        });

        wait_for_listener(&proxy_addr).await;

        let mut stream = TcpStream::connect(&proxy_addr).await.unwrap();
        stream.write_all(b"plain-smoke").await.unwrap();
        stream.shutdown().await.unwrap();

        let mut response = Vec::new();
        stream.read_to_end(&mut response).await.unwrap();
        assert_eq!(response, b"plain-smoke");

        proxy.abort();
    }

    #[tokio::test]
    async fn tls_proxy_forwards_traffic() {
        let backend_addr = spawn_echo_server().await;
        let proxy_addr = unused_loopback_addr();
        let server_config = build_server_config(&tls_config("tls"), "tls").unwrap();

        let proxy = tokio::spawn({
            let proxy_addr = proxy_addr.clone();
            let backend_addr = backend_addr.clone();
            async move {
                run_tls_proxy(&proxy_addr, &backend_addr, server_config)
                    .await
                    .unwrap();
            }
        });

        wait_for_listener(&proxy_addr).await;

        let client_config = ClientConfig::builder()
            .with_root_certificates(root_store())
            .with_no_client_auth();
        let connector = TlsConnector::from(Arc::new(client_config));
        let server_name = ServerName::try_from("localhost").unwrap();

        let stream = TcpStream::connect(&proxy_addr).await.unwrap();
        let mut tls_stream = connector.connect(server_name, stream).await.unwrap();
        tls_stream.write_all(b"tls-smoke").await.unwrap();
        tls_stream.shutdown().await.unwrap();

        let mut response = vec![0u8; b"tls-smoke".len()];
        tls_stream.read_exact(&mut response).await.unwrap();
        assert_eq!(response, b"tls-smoke");

        proxy.abort();
    }

    #[tokio::test]
    async fn mtls_proxy_requires_client_certificate_and_forwards_traffic() {
        let backend_addr = spawn_echo_server().await;
        let proxy_addr = unused_loopback_addr();
        let server_config = build_server_config(&tls_config("mtls"), "mtls").unwrap();

        let proxy = tokio::spawn({
            let proxy_addr = proxy_addr.clone();
            let backend_addr = backend_addr.clone();
            async move {
                run_tls_proxy(&proxy_addr, &backend_addr, server_config)
                    .await
                    .unwrap();
            }
        });

        wait_for_listener(&proxy_addr).await;

        let no_client_auth = ClientConfig::builder()
            .with_root_certificates(root_store())
            .with_no_client_auth();
        let no_cert_connector = TlsConnector::from(Arc::new(no_client_auth));
        let server_name = ServerName::try_from("localhost").unwrap();
        let stream = TcpStream::connect(&proxy_addr).await.unwrap();
        if let Ok(mut tls_stream) = no_cert_connector.connect(server_name.clone(), stream).await {
            let write_result = tls_stream.write_all(b"blocked").await;
            if write_result.is_ok() {
                let mut buf = [0u8; 1];
                let read_result =
                    tokio::time::timeout(Duration::from_secs(1), tls_stream.read(&mut buf)).await;
                match read_result {
                    Ok(Ok(0)) | Ok(Err(_)) | Err(_) => {}
                    Ok(Ok(_)) => panic!("mTLS mode accepted traffic without a client certificate"),
                }
            }
        }

        let client_config = ClientConfig::builder()
            .with_root_certificates(root_store())
            .with_client_auth_cert(
                load_certs(&test_cert_path("client.pem")).unwrap(),
                load_private_key(&test_cert_path("client.key")).unwrap(),
            )
            .unwrap();
        let connector = TlsConnector::from(Arc::new(client_config));

        let stream = TcpStream::connect(&proxy_addr).await.unwrap();
        let mut tls_stream = connector.connect(server_name, stream).await.unwrap();
        tls_stream.write_all(b"mtls-smoke").await.unwrap();
        tls_stream.shutdown().await.unwrap();

        let mut response = vec![0u8; b"mtls-smoke".len()];
        tls_stream.read_exact(&mut response).await.unwrap();
        assert_eq!(response, b"mtls-smoke");

        proxy.abort();
    }
}
