#![forbid(unsafe_code)]

use std::{path::PathBuf, time::Duration};

use anyhow::{Context, Result, bail};
use gstreamer as gst;
use gstreamer::prelude::*;
use gstreamer_app as gst_app;
use lesavka_common::lesavka::{
    AudioEncoding, AudioPacket, UpstreamMediaBundle, VideoPacket, relay_client::RelayClient,
};
use tokio::sync::mpsc;
use tokio_stream::wrappers::ReceiverStream;
use tonic::Request;
use tonic::transport::{Certificate, Channel, ClientTlsConfig, Identity};

const DEFAULT_SERVER: &str = "http://127.0.0.1:50051";
const DEFAULT_SAMPLE_RATE: u32 = 48_000;
const DEFAULT_CHANNELS: u32 = 2;
const DEFAULT_JPEG_QUALITY: i32 = 82;
const MARKER_BITS: usize = 32;
const MARKER_COLUMNS: usize = 16;
const BAND_MARKER_ROWS: usize = 64;
const BAND_MARKER_BITS: usize = 32;
const BAND_MARKER_CELL: usize = 4;

#[derive(Clone, Debug)]
struct Args {
    server: String,
    width: usize,
    height: usize,
    fps: u32,
    duration: Duration,
    jpeg_quality: i32,
    session_id: u64,
    artifact_dir: Option<PathBuf>,
    print_every: u64,
    max_frame_bytes: usize,
    tls_ca: Option<PathBuf>,
    tls_client_cert: Option<PathBuf>,
    tls_client_key: Option<PathBuf>,
    tls_domain: Option<String>,
}

impl Args {
    fn parse() -> Result<Self> {
        let mut args = Self {
            server: DEFAULT_SERVER.to_string(),
            width: 1280,
            height: 720,
            fps: 30,
            duration: Duration::from_secs(300),
            jpeg_quality: DEFAULT_JPEG_QUALITY,
            session_id: unix_millis(),
            artifact_dir: None,
            print_every: 150,
            max_frame_bytes: env_usize("LESAVKA_SYNTHETIC_MAX_FRAME_BYTES").unwrap_or(0),
            tls_ca: env_path("LESAVKA_TLS_CA").or_else(|| default_pki_path("ca.crt")),
            tls_client_cert: env_path("LESAVKA_TLS_CLIENT_CERT")
                .or_else(|| default_pki_path("client.crt")),
            tls_client_key: env_path("LESAVKA_TLS_CLIENT_KEY")
                .or_else(|| default_pki_path("client.key")),
            tls_domain: std::env::var("LESAVKA_TLS_DOMAIN")
                .ok()
                .filter(|value| !value.trim().is_empty()),
        };
        let mut it = std::env::args().skip(1);
        while let Some(flag) = it.next() {
            match flag.as_str() {
                "--server" => args.server = next_value(&mut it, &flag)?,
                "--width" => args.width = parse_next(&mut it, &flag)?,
                "--height" => args.height = parse_next(&mut it, &flag)?,
                "--fps" => args.fps = parse_next(&mut it, &flag)?,
                "--duration" => {
                    let seconds: f64 = parse_next(&mut it, &flag)?;
                    args.duration = Duration::from_secs_f64(seconds.max(0.0));
                }
                "--jpeg-quality" => args.jpeg_quality = parse_next(&mut it, &flag)?,
                "--session-id" => args.session_id = parse_next(&mut it, &flag)?,
                "--artifact-dir" => {
                    args.artifact_dir = Some(PathBuf::from(next_value(&mut it, &flag)?))
                }
                "--print-every" => args.print_every = parse_next(&mut it, &flag)?,
                "--max-frame-bytes" => args.max_frame_bytes = parse_next(&mut it, &flag)?,
                "--tls-ca" => args.tls_ca = Some(PathBuf::from(next_value(&mut it, &flag)?)),
                "--tls-client-cert" => {
                    args.tls_client_cert = Some(PathBuf::from(next_value(&mut it, &flag)?))
                }
                "--tls-client-key" => {
                    args.tls_client_key = Some(PathBuf::from(next_value(&mut it, &flag)?))
                }
                "--tls-domain" => args.tls_domain = Some(next_value(&mut it, &flag)?),
                "--mode" => {
                    let value = next_value(&mut it, &flag)?;
                    let (width, height, fps) = parse_mode(&value)?;
                    args.width = width;
                    args.height = height;
                    args.fps = fps;
                }
                "--help" | "-h" => {
                    print_help();
                    std::process::exit(0);
                }
                other => bail!("unknown argument {other:?}; pass --help for usage"),
            }
        }
        if args.width == 0 || args.height == 0 || args.fps == 0 {
            bail!("width, height, and fps must be positive");
        }
        args.jpeg_quality = args.jpeg_quality.clamp(1, 100);
        Ok(args)
    }

    fn frame_step_us(&self) -> u64 {
        (1_000_000_u64 / u64::from(self.fps)).max(1)
    }

    fn total_frames(&self) -> u64 {
        let frames = self.duration.as_secs_f64() * f64::from(self.fps);
        frames.ceil().max(1.0) as u64
    }
}

struct MjpegEncoder {
    src: gst_app::AppSrc,
    sink: gst_app::AppSink,
    pipeline: gst::Pipeline,
    width: usize,
    height: usize,
    frame_step_us: u64,
}

#[derive(Clone, Copy, Debug, Default)]
struct EncodeStats {
    frames: u64,
    sent_frames: u64,
    total_bytes: u128,
    min_bytes: usize,
    max_bytes: usize,
    oversize_frames: u64,
}

impl EncodeStats {
    fn record(&mut self, bytes: usize, max_frame_bytes: usize) {
        self.frames = self.frames.saturating_add(1);
        self.total_bytes = self.total_bytes.saturating_add(bytes as u128);
        self.min_bytes = if self.frames == 1 {
            bytes
        } else {
            self.min_bytes.min(bytes)
        };
        self.max_bytes = self.max_bytes.max(bytes);
        if max_frame_bytes > 0 && bytes > max_frame_bytes {
            self.oversize_frames = self.oversize_frames.saturating_add(1);
        }
    }

    fn record_sent(&mut self) {
        self.sent_frames = self.sent_frames.saturating_add(1);
    }

    fn mean_bytes(&self) -> usize {
        if self.frames == 0 {
            0
        } else {
            (self.total_bytes / u128::from(self.frames)).min(usize::MAX as u128) as usize
        }
    }
}

impl MjpegEncoder {
    fn new(args: &Args) -> Result<Self> {
        gst::init().context("gst init")?;
        let width = args.width as i32;
        let height = args.height as i32;
        let fps = args.fps as i32;
        let raw_caps = gst::Caps::builder("video/x-raw")
            .field("format", "RGB")
            .field("width", width)
            .field("height", height)
            .field("framerate", gst::Fraction::new(fps, 1))
            .build();
        let jpeg_caps = gst::Caps::builder("image/jpeg")
            .field("parsed", true)
            .field("width", width)
            .field("height", height)
            .field("framerate", gst::Fraction::new(fps, 1))
            .build();
        let pipeline = gst::Pipeline::new();
        let src = gst::ElementFactory::make("appsrc")
            .name("lesavka_synthetic_uplink_src")
            .build()?
            .downcast::<gst_app::AppSrc>()
            .expect("appsrc");
        src.set_is_live(false);
        src.set_format(gst::Format::Time);
        src.set_property("do-timestamp", false);
        src.set_caps(Some(&raw_caps));
        let convert = gst::ElementFactory::make("videoconvert").build()?;
        let encoder = gst::ElementFactory::make("jpegenc")
            .property("quality", args.jpeg_quality)
            .build()?;
        let capsfilter = gst::ElementFactory::make("capsfilter")
            .property("caps", &jpeg_caps)
            .build()?;
        let sink = gst::ElementFactory::make("appsink")
            .name("lesavka_synthetic_uplink_sink")
            .property("sync", false)
            .property("emit-signals", false)
            .property("max-buffers", 8u32)
            .build()?
            .downcast::<gst_app::AppSink>()
            .expect("appsink");
        pipeline.add_many([
            src.upcast_ref(),
            &convert,
            &encoder,
            &capsfilter,
            sink.upcast_ref(),
        ])?;
        gst::Element::link_many([
            src.upcast_ref(),
            &convert,
            &encoder,
            &capsfilter,
            sink.upcast_ref(),
        ])?;
        pipeline
            .set_state(gst::State::Playing)
            .context("starting synthetic MJPEG encoder")?;
        Ok(Self {
            src,
            sink,
            pipeline,
            width: args.width,
            height: args.height,
            frame_step_us: args.frame_step_us(),
        })
    }

    fn encode(&mut self, sequence: u64) -> Result<Vec<u8>> {
        let pts_us = sequence.saturating_mul(self.frame_step_us);
        let mut buffer =
            gst::Buffer::from_slice(synthetic_rgb_frame(self.width, self.height, sequence));
        if let Some(meta) = buffer.get_mut() {
            let pts = gst::ClockTime::from_useconds(pts_us);
            meta.set_pts(Some(pts));
            meta.set_dts(Some(pts));
            meta.set_duration(Some(gst::ClockTime::from_useconds(self.frame_step_us)));
        }
        self.src
            .push_buffer(buffer)
            .context("encoding synthetic frame")?;
        let sample = self
            .sink
            .pull_sample()
            .context("pulling encoded synthetic frame")?;
        let buffer = sample
            .buffer()
            .context("encoded synthetic frame had no buffer")?;
        let map = buffer
            .map_readable()
            .context("mapping encoded synthetic frame")?;
        add_jpeg_app4_integrity(map.as_slice(), sequence)
    }
}

impl Drop for MjpegEncoder {
    fn drop(&mut self) {
        let _ = self.src.end_of_stream();
        let _ = self.pipeline.set_state(gst::State::Null);
    }
}

async fn connect_channel(args: &Args) -> Result<Channel> {
    let mut endpoint =
        tonic::transport::Channel::from_shared(args.server.clone())?.tcp_nodelay(true);
    if is_https(&args.server) {
        endpoint = endpoint
            .tls_config(client_tls_config(args)?)
            .context("configuring synthetic uplink TLS")?;
    }
    endpoint
        .connect()
        .await
        .with_context(|| format!("connecting to {}", args.server))
}

fn client_tls_config(args: &Args) -> Result<ClientTlsConfig> {
    let mut tls = ClientTlsConfig::new().domain_name(
        args.tls_domain
            .clone()
            .or_else(|| host_from_uri(&args.server))
            .unwrap_or_else(|| "lesavka-server".to_string()),
    );
    let ca_path = args
        .tls_ca
        .as_ref()
        .context("https synthetic uplink requires --tls-ca or LESAVKA_TLS_CA")?;
    let cert_path = args
        .tls_client_cert
        .as_ref()
        .context("https synthetic uplink requires --tls-client-cert or LESAVKA_TLS_CLIENT_CERT")?;
    let key_path = args
        .tls_client_key
        .as_ref()
        .context("https synthetic uplink requires --tls-client-key or LESAVKA_TLS_CLIENT_KEY")?;
    let ca = std::fs::read(ca_path)
        .with_context(|| format!("reading TLS CA certificate {}", ca_path.display()))?;
    tls = tls.ca_certificate(Certificate::from_pem(ca));
    let cert = std::fs::read(cert_path)
        .with_context(|| format!("reading TLS client certificate {}", cert_path.display()))?;
    let key = std::fs::read(key_path)
        .with_context(|| format!("reading TLS client key {}", key_path.display()))?;
    Ok(tls.identity(Identity::from_pem(cert, key)))
}

fn is_https(server: &str) -> bool {
    server.trim_start().starts_with("https://")
}

fn host_from_uri(server: &str) -> Option<String> {
    let rest = server.split_once("://")?.1;
    let host_port = rest.split('/').next().unwrap_or(rest);
    let host = host_port
        .rsplit_once('@')
        .map(|(_, host)| host)
        .unwrap_or(host_port);
    if host.starts_with('[') {
        return host
            .split_once(']')
            .map(|(value, _)| value.trim_start_matches('[').to_string());
    }
    Some(host.split(':').next().unwrap_or(host).to_string()).filter(|host| !host.is_empty())
}

fn env_path(name: &str) -> Option<PathBuf> {
    std::env::var_os(name)
        .filter(|value| !value.is_empty())
        .map(PathBuf::from)
}

fn default_pki_path(file_name: &str) -> Option<PathBuf> {
    let home = std::env::var_os("HOME")?;
    Some(
        PathBuf::from(home)
            .join(".config")
            .join("lesavka")
            .join("pki")
            .join(file_name),
    )
}

#[tokio::main(flavor = "current_thread")]
async fn main() -> Result<()> {
    let args = Args::parse()?;
    if let Some(dir) = &args.artifact_dir {
        std::fs::create_dir_all(dir).with_context(|| format!("creating {}", dir.display()))?;
        std::fs::write(
            dir.join("command.txt"),
            std::env::args().collect::<Vec<_>>().join(" ") + "\n",
        )?;
        write_summary(&args, None, None)?;
    }

    let channel = connect_channel(&args).await?;
    let mut client = RelayClient::new(channel);
    let (tx, rx) = mpsc::channel::<UpstreamMediaBundle>(8);
    let response_task = tokio::spawn(async move {
        let response = client
            .stream_webcam_media(Request::new(ReceiverStream::new(rx)))
            .await
            .context("opening StreamWebcamMedia")?;
        let mut inbound = response.into_inner();
        while inbound
            .message()
            .await
            .context("reading StreamWebcamMedia response")?
            .is_some()
        {}
        Ok::<(), anyhow::Error>(())
    });

    let mut encoder = MjpegEncoder::new(&args)?;
    let mut encode_stats = EncodeStats::default();
    let frame_step = Duration::from_micros(args.frame_step_us());
    let started = tokio::time::Instant::now() + Duration::from_millis(250);
    let total_frames = args.total_frames();
    eprintln!(
        "lesavka synthetic uplink: mode={}x{}@{} frames={} server={} session={}",
        args.width, args.height, args.fps, total_frames, args.server, args.session_id
    );
    for sequence in 0..total_frames {
        tokio::time::sleep_until(started + duration_mul(frame_step, sequence)).await;
        let pts_us = sequence.saturating_mul(args.frame_step_us());
        let encoded = encoder.encode(sequence)?;
        encode_stats.record(encoded.len(), args.max_frame_bytes);
        if args.max_frame_bytes > 0 && encoded.len() > args.max_frame_bytes {
            write_summary(
                &args,
                Some(&encode_stats),
                Some("encoded_frame_over_budget"),
            )?;
            bail!(
                "encoded synthetic frame {sequence} is {} bytes, above --max-frame-bytes {}; lower --jpeg-quality or use a more compressible synthetic pattern before trusting the RCT probe",
                encoded.len(),
                args.max_frame_bytes
            );
        }
        let bundle = synthetic_bundle(&args, sequence, pts_us, encoded);
        if tx.send(bundle).await.is_err() {
            let closed_reason = format!(
                "StreamWebcamMedia closed before accepting synthetic frame {sequence}; disconnect or pause any live Lesavka client upstream before running the isolated RCT probe"
            );
            let response_result = response_task
                .await
                .context("joining StreamWebcamMedia task after early close")?;
            match response_result {
                Ok(()) => {
                    write_summary(&args, Some(&encode_stats), Some(&closed_reason))?;
                    bail!("{closed_reason}");
                }
                Err(err) => {
                    let error_reason = format!(
                        "StreamWebcamMedia closed before accepting synthetic frame {sequence}: {err:#}"
                    );
                    write_summary(&args, Some(&encode_stats), Some(&error_reason))?;
                    return Err(err)
                        .context("StreamWebcamMedia closed before accepting synthetic frame");
                }
            }
        }
        encode_stats.record_sent();
        if args.print_every > 0 && sequence > 0 && sequence % args.print_every == 0 {
            eprintln!("sent synthetic frame {sequence}/{total_frames}");
        }
    }
    drop(tx);
    response_task
        .await
        .context("joining StreamWebcamMedia task")??;
    write_summary(&args, Some(&encode_stats), Some("complete"))?;
    eprintln!("lesavka synthetic uplink complete: frames={total_frames}");
    Ok(())
}

include!("lesavka_synthetic_uplink/support.rs");

#[cfg(test)]
#[path = "lesavka_synthetic_uplink/tests/mod.rs"]
mod tests;
