use anyhow_ext::{Context, Result, bail};
use chrono::{Duration as ChronoDuration, Local};
use clap::{ArgGroup, Parser};
use ssh2::Session;
use std::io::Read;
use std::process::exit;
use std::{path::PathBuf, time::Duration as StdDuration};
use tokio::fs::{File, OpenOptions, create_dir_all, read_to_string};
use tokio::io::AsyncWriteExt;
use tokio::sync::mpsc;
use tokio::time::{Duration as TokioDuration, sleep};
use tokio::{
    io::{AsyncBufReadExt, BufReader},
    net::TcpStream as TokioTcpStream,
};

const BANNER: &str = r#"
// ┌─────────────────────────────────────────────────────────────────────┐
// │ Matrix NARUTO                                                       │
// │                                                                     │
// │    1        0        1        0        1        2                   │
// │    0        9        0        1        6        1                   │
// │    1        0        8        1        0        0                   │
// │    0        1        0        9        1        1                   │
// │    1        0        1        0        0        0                   │
// │    0        1        0        1        0        5                   │
// │                                                                     │
// │  混乱并不是错误，它只是未被理解的秩序   *                               │
// └─────────────────────────────────────────────────────────────────────┘
"#;


#[derive(Parser, Debug)]
#[command(
    author,
    version,
    about = "Joker logs Auto Save Tools, By Bruce Shi",
    group(
        ArgGroup::new("mode")
            .required(true)
            .args(["standard", "pro"])
    )
)]

struct Cli {
    /// 标准基站
    #[arg(short = 's', long = "standard")]
    standard: bool,

    /// P基站
    #[arg(short = 'p', long = "pro")]
    pro: bool,

    /// 包含 IP 的 txt，每行一个
    #[arg(value_name = "IP_FILE")]
    ip_file: String,

    /// SSH 远端日志路径
    #[arg(long, default_value = "/var/log/messages", hide = true)]
    remote: String,

    /// SSH 端口
    #[arg(long, default_value_t = 22, hide = true)]
    port: u16,

    /// SSH 掉线重连间隔秒
    #[arg(long, default_value_t = 3, hide = true)]
    reconnect_delay: u64,
}

#[tokio::main]
async fn main() -> Result<()> {
    let cli = Cli::parse();
    print!("{}", BANNER);
    let ip_file = cli.ip_file;

    ctrlc::set_handler(|| {
        println!("get ctrl, not exit it ");
        exit(0)
    })
    .expect("ctrl c exit failed");

    if cli.pro {
        exec_pro(&ip_file).await?;
    }

    if cli.standard {
        exec_standard(&ip_file).await?;
    }

    Ok(())
}

/// 执行标准基站获取日志流程
async fn exec_standard(files: &str) -> Result<()> {
    let ip_file = files;
    let content = read_to_string(ip_file).await?;
    let ips: Vec<String> = content
        .lines()
        .filter(|s| !s.trim().is_empty())
        .map(|s| s.trim().to_string())
        .collect();

    if ips.is_empty() {
        bail!("IP file is empty!");
    }
    println!("Loaded {} IPs.", ips.len());

    // 创建日志目录
    create_dir_all(LOG_DIR).await?;

    // 为每个 IP 启动任务
    for ip in ips {
        tokio::spawn(async move {
            if let Err(e) = handle_ip(ip.clone()).await {
                eprintln!("Task for IP {} failed: {:?}", ip, e);
            }
        });
    }

    // 防止主线程退出
    loop {
        sleep(TokioDuration::from_secs(1800)).await;
    }
}

/// 单 IP 异步处理：读取 + 写入队列
async fn handle_ip(ip: String) -> Result<()> {
    println!("Connecting to {}:{} ...", ip, EW_LOG_PORT);

    let ip_port = format!("{}:{}", ip, EW_LOG_PORT);
    let stream = TokioTcpStream::connect(ip_port).await?;

    let reader = BufReader::new(stream);
    let mut lines = reader.lines();

    // 创建消息队列（一次最多缓存 10k 行）
    let (tx, rx) = mpsc::channel::<String>(10_000);

    let ip_for_write = ip.clone();
    println!("Connected to {}:{} success", ip, EW_LOG_PORT);
    // 启动写文件任务（独立）
    tokio::spawn(async move {
        if let Err(e) = log_writer(ip_for_write.clone(), rx).await {
            eprintln!("({}) log write error: {:?}", ip_for_write, e);
        }
    });
    // telnet 输入 → 写入队列
    while let Ok(Some(line)) = lines.next_line().await {
        if tx.send(line).await.is_err() {
            println!("({}) writer closed.", ip);
            break;
        }
    }

    println!("({}) Telnet disconnected", ip);
    Ok(())
}

/// 处理写文件任务（批量写 + 自动滚动）
async fn log_writer(ip: String, mut rx: mpsc::Receiver<String>) -> Result<()> {
    let mut file = create_new_log_file(&ip).await?;
    let mut file_time = Local::now();

    // 64KB 缓冲区（高速）
    let mut buffer = String::with_capacity(64 * 1024);
    loop {
        tokio::select! {
            maybe_line = rx.recv() => {
                match maybe_line {
                    Some(line) => {
                        buffer.push_str(&line);
                        buffer.push('\n');
                        // println!("recv {}",buffer);
                        if buffer.len() >= 64 * 1024 {
                            file.write_all(buffer.as_bytes()).await?;
                            buffer.clear();

                            // 文件滚动：超过 512MB
                            let metadata = file.metadata().await?;
                            if metadata.len() > MAX_FILE_SIZE {
                                println!("({}) Log file > 512MB, rotate...", ip);
                                file = create_new_log_file(&ip).await?;
                                file_time = Local::now();
                            }
                        }
                    }
                    None => {
                        // 发送端关闭，刷掉残留并退出
                        if !buffer.is_empty() {
                            file.write_all(buffer.as_bytes()).await?;
                            buffer.clear();
                        }
                        file.flush().await?;
                        break;
                    }
                }
            }

            // 每秒 flush 一次
            _ = sleep(TokioDuration::from_secs(1)) => {
                if !buffer.is_empty() {
                    file.write_all(buffer.as_bytes()).await?;
                    buffer.clear();
                }

                // 按天滚动
                if Local::now() - file_time > ChronoDuration::days(1) {
                    println!("({}) Log file > 1 day, rotate...", ip);
                    file = create_new_log_file(&ip).await?;
                    file_time = Local::now();
                }
            }
        }
    }

    Ok(())
}

/// 创建新的日志文件：格式为 `IP_YYYYMMDD_HHMMSS.log`
async fn create_new_log_file(ip: &str) -> Result<File> {
    let now = Local::now();
    let filename = format!(
        "{}_{}.log",
        ip.replace(":", "_"),
        now.format("%Y%m%d_%H%M%S")
    );

    let mut path = PathBuf::from(LOG_DIR);
    path.push(filename);

    println!("Creating new log file: {:?}", path);

    let file = OpenOptions::new()
        .create(true)
        .append(true)
        .open(path)
        .await?;

    Ok(file)
}

// ----------------------Pro基站处理逻辑---------------------------------- //

fn connect_ssh_sync(ip: &str) -> Result<Session> {
    let addr = format!("{}:22", ip);
    let tcp = std::net::TcpStream::connect(&addr)
        .with_context(|| format!("connect tcp {} error", addr))?;

    tcp.set_read_timeout(Some(StdDuration::from_secs(30))).ok();
    tcp.set_write_timeout(Some(StdDuration::from_secs(30))).ok();

    let mut sess = Session::new().context("create ssh session")?;
    sess.set_tcp_stream(tcp);
    sess.handshake().context("ssh handshake")?;
    sess.userauth_password(PRO_SSH_USER, PRO_SSH_PASSWORD)
        .context("auth by password")?;

    if !sess.authenticated() {
        bail!("ssh not authenticated");
    }
    Ok(sess)
}

// 简单 shell escape：把路径用单引号包起来并处理单引号
fn shell_escape(s: &str) -> String {
    if s.is_empty() {
        "''".to_string()
    } else if s
        .bytes()
        .all(|c| c.is_ascii_alphanumeric() || b"/._-".contains(&c))
    {
        s.to_string()
    } else {
        let escaped = s.replace('\'', r"'\''");
        format!("'{}'", escaped)
    }
}

async fn exec_pro(ipfile: &str) -> Result<()> {
    let content = read_to_string(ipfile).await?;
    let ips: Vec<String> = content
        .lines()
        .filter(|s| !s.trim().is_empty())
        .map(|s| s.trim().to_string())
        .collect();

    if ips.is_empty() {
        bail!("IP file is empty!");
    }

    // 默认参数
    let reconnect_delay: u64 = 3;
    let tail_lines: u32 = 100;
    let remote = "/var/log/messages";
    create_dir_all(LOG_DIR).await?;

    for ip in ips {
        let ip_clone = ip.clone();
        let remote = remote.to_string();
        tokio::spawn(async move {
            if let Err(e) =
                tail_pro_host(ip_clone.clone(), &remote, tail_lines, reconnect_delay).await
            {
                eprintln!("Task for IP {} failed: {:?}", ip_clone, e);
            }
        });
    }

    loop {
        sleep(TokioDuration::from_secs(1800)).await;
    }
}

async fn tail_pro_host(
    ip: String,
    remote: &str,
    tail_lines: u32,
    reconnect_delay: u64,
) -> Result<()> {
    let remote = remote.to_string();

    loop {
        // pipeline：ssh(blocking) -> tx -> log_writer(async)
        let (tx, rx) = mpsc::channel::<String>(10_000);

        let ip_for_write = ip.clone();
        tokio::spawn(async move {
            if let Err(e) = log_writer(ip_for_write.clone(), rx).await {
                eprintln!("({}) log write error: {:?}", ip_for_write, e);
            }
        });

        eprintln!("[{}] connect -> tail {}", ip, remote);

        let ip_clone = ip.clone();
        let remote_clone = remote.clone();

        // 关键：把 ssh2 的阻塞读，丢到 blocking 线程池
        let join = tokio::task::spawn_blocking(move || -> Result<()> {
            let sess = connect_ssh_sync(&ip_clone)?;

            let mut channel = sess
                .channel_session()
                .with_context(|| format!("channel_session {}", ip_clone))?;

            let cmd = if tail_lines == 0 {
                format!("tail -F {} 2>&1", shell_escape(&remote_clone))
            } else {
                format!(
                    "tail -n {} -F {} 2>&1",
                    tail_lines,
                    shell_escape(&remote_clone)
                )
            };

            eprintln!("[ssh {}] exec: {}", ip_clone, cmd);
            channel
                .exec(&cmd)
                .with_context(|| format!("exec tail {}", ip_clone))?;

            let mut buf = [0u8; 8192];
            loop {
                match channel.read(&mut buf) {
                    Ok(0) => {
                        eprintln!("[ssh {}] remote closed stream", ip_clone);
                        break;
                    }
                    Ok(nread) => {
                        let chunk = String::from_utf8_lossy(&buf[..nread]).to_string();
                        if tx.blocking_send(chunk).is_err() {
                            break;
                        }
                    }
                    Err(e) => {
                        eprintln!("[read {}] {:#}", ip_clone, e);
                        break;
                    }
                }
            }

            let _ = channel.close();
            let _ = channel.wait_close();
            Ok(())
        });

        // 等这次 blocking tail 结束（断线/报错都会结束）
        if let Err(e) = join.await {
            eprintln!("[join {}] {:#}", ip, e);
        }

        // 关闭发送端，保证 writer 能刷盘退出/继续滚动
        // （这里 tx 已经 move 进 blocking 闭包了，随着闭包结束自动 drop）
        std::thread::sleep(StdDuration::from_secs(reconnect_delay));
    }
}

[package]
name = "autosave"
version = "0.1.0"
edition = "2024"


[dependencies]
ssh2 = "0.9"
anyhow_ext = "0.1.0"
clap = { version = "4", features = ["derive"] }
chrono = "0.4"
ctrlc = "3"
tokio = { version = "1", features = ["rt-multi-thread", "macros", "fs", "net", "io-util", "sync", "time"] }

本地测试通过