use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use axum::{body::Body, extract::State, http::{Request, StatusCode}, response::Response};
use reqwest::Client;

pub struct UpstreamNode {
    pub url: String,
    pub api_key: String,
    pub active_conns: AtomicUsize,
    pub is_healthy: std::sync::atomic::AtomicBool,
}

pub struct AppState {
    pub nodes: Vec<Arc<UpstreamNode>>,
    pub rr_index: AtomicUsize,
    pub http_client: Client,
}

impl AppState {
    /// 基于最少活跃连接与轮询的混合调度算法
    pub fn select_node(&self) -> Option<Arc<UpstreamNode>> {
        let len = self.nodes.len();
        if len == 0 {
            return None;
        }
        
        let start = self.rr_index.fetch_add(1, Ordering::Relaxed);
        for i in 0..len {
            let idx = (start + i) % len;
            let node = &self.nodes[idx];
            if node.is_healthy.load(Ordering::Acquire) {
                return Some(Arc::clone(node));
            }
        }
        None
    }
}

/// 核心流式转发 Handler
pub async fn proxy_llm_stream(
    State(state): State<Arc<AppState>>,
    req: Request<Body>,
) -> Result<Response<Body>, StatusCode> {
    let node = state.select_node().ok_or(StatusCode::SERVICE_UNAVAILABLE)?;
    
    // 记录活跃连接数（RAII 模式确保异常退出时也能准确递减）
    node.active_conns.fetch_add(1, Ordering::Relaxed);
    let node_clone = Arc::clone(&node);
    scopeguard::defer! {
        node_clone.active_conns.fetch_sub(1, Ordering::Relaxed);
    }

    // 构造上游请求并直接透传数据流
    let target_url = format!("{}/v1/chat/completions", node.url);
    let upstream_resp = state.http_client
        .post(&target_url)
        .bearer_auth(&node.api_key)
        .body(reqwest::Body::wrap_stream(req.into_body().into_data_stream()))
        .send()
        .await
        .map_err(|_| {
            // 触发被动熔断标记
            node.is_healthy.store(false, Ordering::Release);
            StatusCode::BAD_GATEWAY
        })?;

    // 将上游的字节流零拷贝桥接回客户端
    let mut response_builder = Response::builder().status(upstream_resp.status());
    if let Some(headers) = response_builder.headers_mut() {
        *headers = upstream_resp.headers().clone();
    }

    let stream = upstream_resp.bytes_stream();
    response_builder
        .body(Body::from_stream(stream))
        .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)
}

FROM rust:1.82-alpine AS builder
RUN apk add --no-cache musl-dev pkgconfig openssl-dev openssl-libs-static
WORKDIR /app
COPY . .
RUN cargo build --release --target x86_64-unknown-linux-musl

FROM scratch
COPY --from=builder /app/target/x86_64-unknown-linux-musl/release/llm-gateway /llm-gateway
EXPOSE 8080
ENTRYPOINT ["/llm-gateway"]

# 开启 BBR 拥塞控制算法，显著改善跨公网 SSE 流式传输的尾部延迟与抗丢包能力
net.core.default_qdisc = fq
net.ipv4.tcp_congestion_control = bbr

# 扩大本地临时端口范围，支撑高并发上游建连
net.ipv4.ip_local_port_range = 10240 65535

# 开启 TIME_WAIT 快速复用并增大最大文件描述符限制
net.ipv4.tcp_tw_reuse = 1
net.core.somaxconn = 65535
net.ipv4.tcp_max_syn_backlog = 65535
fs.file-max = 2097152