mod receiver;
mod s3;
mod watcher;

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

use anyhow::{Context, Result};
use axum::extract::State;
use axum::http::StatusCode;
use axum::response::Json;
use axum::routing::get;
use axum::Router;
use serde::{Deserialize, Serialize};
use tokio::sync::RwLock;
use tower_http::cors::CorsLayer;
use tower_http::services::ServeDir;
use tower_http::trace::TraceLayer;
use tracing::info;

use receiver::ReceiverHandle;
use s3::S3Uploader;
use watcher::SegmentWatcher;

#[derive(Clone)]
struct AppState {
    receiver: Arc<RwLock<Option<ReceiverHandle>>>,
    uploader: Option<Arc<S3Uploader>>,
    config: Arc<Config>,
    stats: Arc<RwLock<Stats>>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
struct Config {
    /// Port to listen for SRT stream
    srt_port: u16,
    /// Directory to store segments
    segment_dir: PathBuf,
    /// Segment duration in seconds
    segment_duration: u32,
    /// S3 bucket name
    s3_bucket: String,
    /// S3 prefix for uploads
    s3_prefix: String,
    /// HTTP API port
    api_port: u16,
    /// Delete local files after S3 upload
    delete_after_upload: bool,
    /// ffmpeg path
    ffmpeg_path: PathBuf,
    /// YouTube RTMP URL (rtmp://a.rtmp.youtube.com/live2)
    youtube_rtmp_url: String,
    /// YouTube stream key
    youtube_stream_key: String,
    /// Enable YouTube streaming
    youtube_enabled: bool,
    /// HLS output directory (for web playback)
    hls_dir: PathBuf,
    /// Enable HLS output for web playback
    hls_enabled: bool,
    /// HLS segment duration in seconds
    hls_segment_duration: u32,
}

impl Default for Config {
    fn default() -> Self {
        Self {
            srt_port: 6000,
            segment_dir: PathBuf::from("/root/stream/segments"),
            segment_duration: 60,
            s3_bucket: String::new(),
            s3_prefix: "stream".into(),
            api_port: 8080,
            delete_after_upload: true,
            ffmpeg_path: PathBuf::from("ffmpeg"),
            youtube_rtmp_url: "rtmp://a.rtmp.youtube.com/live2".into(),
            youtube_stream_key: String::new(),
            youtube_enabled: false,
            hls_dir: PathBuf::from("/root/stream/hls"),
            hls_enabled: true, // Enable HLS by default for web playback
            hls_segment_duration: 4,
        }
    }
}

#[derive(Debug, Default, Serialize)]
struct Stats {
    receiving: bool,
    segments_uploaded: u64,
    bytes_uploaded: u64,
    last_segment: Option<String>,
    errors: Vec<String>,
}

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

    let config = load_config()?;
    info!("Starting stream server with config: {:?}", config);

    // Ensure segment directory exists
    tokio::fs::create_dir_all(&config.segment_dir)
        .await
        .with_context(|| format!("create segment dir {}", config.segment_dir.display()))?;

    // Initialize S3 uploader only when configured
    let uploader = if config.s3_bucket.trim().is_empty() {
        info!("S3 upload disabled (S3_BUCKET not set)");
        None
    } else {
        let uploader = S3Uploader::new(&config.s3_bucket, &config.s3_prefix).await?;
        Some(Arc::new(uploader))
    };

    let state = AppState {
        receiver: Arc::new(RwLock::new(None)),
        uploader: uploader.clone(),
        config: Arc::new(config.clone()),
        stats: Arc::new(RwLock::new(Stats::default())),
    };

    // Start segment watcher only when S3 upload is enabled
    if state.uploader.is_some() {
        let watcher_state = state.clone();
        tokio::spawn(async move {
            if let Err(e) = run_watcher(watcher_state).await {
                tracing::error!("Watcher error: {e}");
            }
        });
    }

    // Auto-start receiver
    let receiver_state = state.clone();
    tokio::spawn(async move {
        if let Err(e) = start_receiver_internal(&receiver_state).await {
            tracing::error!("Failed to start receiver: {e}");
        }
    });

    // Ensure HLS directory exists
    if config.hls_enabled {
        tokio::fs::create_dir_all(&config.hls_dir)
            .await
            .with_context(|| format!("create HLS dir {}", config.hls_dir.display()))?;
    }

    // Build HTTP API
    let mut app = Router::new()
        .route("/", get(index))
        .route("/status", get(status))
        .route("/start", get(start_receiver))
        .route("/stop", get(stop_receiver))
        .route("/health", get(health));

    // Serve HLS files if enabled
    if config.hls_enabled {
        info!("Serving HLS files from {} at /hls/", config.hls_dir.display());
        app = app.nest_service("/hls", ServeDir::new(&config.hls_dir));
    }

    let app = app
        .layer(TraceLayer::new_for_http())
        .layer(CorsLayer::permissive())
        .with_state(state);

    let addr = SocketAddr::from(([0, 0, 0, 0], config.api_port));
    info!("HTTP API listening on {}", addr);

    let listener = tokio::net::TcpListener::bind(addr).await?;
    axum::serve(listener, app).await?;

    Ok(())
}

fn load_config() -> Result<Config> {
    let mut config = Config::default();

    // Load from environment
    if let Ok(port) = std::env::var("SRT_PORT") {
        config.srt_port = port.parse().context("parse SRT_PORT")?;
    }
    if let Ok(dir) = std::env::var("SEGMENT_DIR") {
        config.segment_dir = PathBuf::from(dir);
    }
    if let Ok(duration) = std::env::var("SEGMENT_DURATION") {
        config.segment_duration = duration.parse().context("parse SEGMENT_DURATION")?;
    }
    if let Ok(bucket) = std::env::var("S3_BUCKET") {
        config.s3_bucket = bucket;
    }
    if let Ok(prefix) = std::env::var("S3_PREFIX") {
        config.s3_prefix = prefix;
    }
    if let Ok(port) = std::env::var("API_PORT") {
        config.api_port = port.parse().context("parse API_PORT")?;
    }
    if let Ok(delete) = std::env::var("DELETE_AFTER_UPLOAD") {
        config.delete_after_upload = delete == "true" || delete == "1";
    }
    if let Ok(path) = std::env::var("FFMPEG_PATH") {
        config.ffmpeg_path = PathBuf::from(path);
    }
    if let Ok(url) = std::env::var("YOUTUBE_RTMP_URL") {
        config.youtube_rtmp_url = url;
    }
    if let Ok(key) = std::env::var("YOUTUBE_STREAM_KEY") {
        config.youtube_stream_key = key;
        config.youtube_enabled = true;
    }
    if let Ok(enabled) = std::env::var("YOUTUBE_ENABLED") {
        config.youtube_enabled = enabled == "true" || enabled == "1";
    }
    if let Ok(dir) = std::env::var("HLS_DIR") {
        config.hls_dir = PathBuf::from(dir);
    }
    if let Ok(enabled) = std::env::var("HLS_ENABLED") {
        config.hls_enabled = enabled == "true" || enabled == "1";
    }
    if let Ok(duration) = std::env::var("HLS_SEGMENT_DURATION") {
        config.hls_segment_duration = duration.parse().context("parse HLS_SEGMENT_DURATION")?;
    }

    Ok(config)
}

async fn run_watcher(state: AppState) -> Result<()> {
    let uploader = match state.uploader.clone() {
        Some(uploader) => uploader,
        None => {
            info!("S3 upload disabled; segment watcher not started");
            return Ok(());
        }
    };
    let mut watcher = SegmentWatcher::new(&state.config.segment_dir)?;

    loop {
        if let Some(segment_path) = watcher.next_segment().await {
            let filename = segment_path
                .file_name()
                .and_then(|n| n.to_str())
                .unwrap_or("unknown")
                .to_string();

            info!("Uploading segment: {}", filename);

            match uploader.upload_file(&segment_path).await {
                Ok(bytes) => {
                    let mut stats = state.stats.write().await;
                    stats.segments_uploaded += 1;
                    stats.bytes_uploaded += bytes;
                    stats.last_segment = Some(filename.clone());

                    if state.config.delete_after_upload {
                        if let Err(e) = tokio::fs::remove_file(&segment_path).await {
                            tracing::warn!("Failed to delete {}: {e}", segment_path.display());
                        } else {
                            info!("Deleted local segment: {}", filename);
                        }
                    }
                }
                Err(e) => {
                    tracing::error!("Upload failed for {}: {e}", filename);
                    let mut stats = state.stats.write().await;
                    stats.errors.push(format!("Upload {}: {e}", filename));
                    if stats.errors.len() > 100 {
                        stats.errors.remove(0);
                    }
                }
            }
        }
    }
}

async fn start_receiver_internal(state: &AppState) -> Result<()> {
    let mut receiver_guard = state.receiver.write().await;
    if receiver_guard.is_some() {
        return Ok(());
    }

    let handle = if state.config.youtube_enabled && !state.config.youtube_stream_key.is_empty() {
        info!("Starting YouTube streaming mode");
        receiver::start_youtube(
            &state.config.ffmpeg_path,
            state.config.srt_port,
            &state.config.youtube_rtmp_url,
            &state.config.youtube_stream_key,
        )
        .await?
    } else if state.config.hls_enabled {
        info!("Starting HLS streaming mode");
        receiver::start_hls(
            &state.config.ffmpeg_path,
            state.config.srt_port,
            &state.config.hls_dir,
            state.config.hls_segment_duration,
        )
        .await?
    } else {
        receiver::start(
            &state.config.ffmpeg_path,
            state.config.srt_port,
            &state.config.segment_dir,
            state.config.segment_duration,
        )
        .await?
    };

    *receiver_guard = Some(handle);
    state.stats.write().await.receiving = true;

    info!(
        "SRT receiver started on port {}",
        state.config.srt_port
    );
    Ok(())
}

async fn index() -> &'static str {
    "Stream Server - POST /start, /stop, GET /status, /health"
}

async fn health() -> StatusCode {
    StatusCode::OK
}

async fn status(State(state): State<AppState>) -> Json<StatusResponse> {
    let stats = state.stats.read().await;
    let receiver = state.receiver.read().await;

    Json(StatusResponse {
        receiving: receiver.is_some(),
        srt_port: state.config.srt_port,
        s3_bucket: state.config.s3_bucket.clone(),
        youtube_enabled: state.config.youtube_enabled,
        youtube_rtmp_url: if state.config.youtube_enabled {
            Some(state.config.youtube_rtmp_url.clone())
        } else {
            None
        },
        hls_enabled: state.config.hls_enabled,
        hls_url: if state.config.hls_enabled {
            Some(format!("/hls/stream.m3u8"))
        } else {
            None
        },
        segments_uploaded: stats.segments_uploaded,
        bytes_uploaded: stats.bytes_uploaded,
        bytes_uploaded_human: human_bytes(stats.bytes_uploaded),
        last_segment: stats.last_segment.clone(),
        recent_errors: stats.errors.iter().rev().take(5).cloned().collect(),
    })
}

#[derive(Serialize)]
struct StatusResponse {
    receiving: bool,
    srt_port: u16,
    s3_bucket: String,
    youtube_enabled: bool,
    youtube_rtmp_url: Option<String>,
    hls_enabled: bool,
    hls_url: Option<String>,
    segments_uploaded: u64,
    bytes_uploaded: u64,
    bytes_uploaded_human: String,
    last_segment: Option<String>,
    recent_errors: Vec<String>,
}

async fn start_receiver(State(state): State<AppState>) -> Result<Json<serde_json::Value>, StatusCode> {
    match start_receiver_internal(&state).await {
        Ok(()) => Ok(Json(serde_json::json!({"status": "started"}))),
        Err(e) => {
            tracing::error!("Failed to start receiver: {e}");
            Err(StatusCode::INTERNAL_SERVER_ERROR)
        }
    }
}

async fn stop_receiver(State(state): State<AppState>) -> Json<serde_json::Value> {
    let mut receiver_guard = state.receiver.write().await;
    if let Some(handle) = receiver_guard.take() {
        handle.stop().await;
        state.stats.write().await.receiving = false;
        info!("SRT receiver stopped");
    }
    Json(serde_json::json!({"status": "stopped"}))
}

fn human_bytes(bytes: u64) -> String {
    const KB: u64 = 1024;
    const MB: u64 = KB * 1024;
    const GB: u64 = MB * 1024;

    if bytes >= GB {
        format!("{:.2} GB", bytes as f64 / GB as f64)
    } else if bytes >= MB {
        format!("{:.2} MB", bytes as f64 / MB as f64)
    } else if bytes >= KB {
        format!("{:.2} KB", bytes as f64 / KB as f64)
    } else {
        format!("{} B", bytes)
    }
}