use std::{error::Error, sync::Arc, time::Duration};

use async_broadcast::broadcast;
use chrono::Utc;
use dashmap::DashMap;
use entity::stream_session;
use rml_rtmp::{
    handshake::{Handshake, HandshakeProcessResult, PeerType},
    sessions::{ServerSession, ServerSessionConfig, ServerSessionEvent, ServerSessionResult},
};
use sea_orm::{DatabaseConnection, IntoActiveModel, sqlx::types::chrono::Local};
use tokio::{
    io::{AsyncReadExt, AsyncWriteExt},
    net::{TcpListener, TcpStream},
    time::{Instant, timeout},
};
use tracing::{debug, info, warn};

use crate::{
    StreamCodec, StreamSession,
    audio::{AACParser, AudioProcesser, OpusAudioFrame},
    codec::{
        CodecParser, VideoFrame, av1::Av1CodecParser, h264::H264CodecParser, h265::H265CodecParser,
    },
};

pub struct Rtmp {
    pub listener: TcpListener,
    pub stream_sessions: Arc<DashMap<i32, StreamSession>>,
    pub db: DatabaseConnection,
}

async fn write_outbound(socket: &mut TcpStream, results: Vec<ServerSessionResult>) {
    for r in results {
        if let ServerSessionResult::OutboundResponse(p) = r {
            socket.write_all(&p.bytes).await.unwrap();
        }
    }
}

impl Rtmp {
    fn parse_video_codec(payload: &[u8]) -> Result<StreamCodec, Box<dyn Error + Send + Sync>> {
        if payload.is_empty() {
            return Err("empty video payload".into());
        }

        let is_ex = payload[0] & 0x80 != 0;

        if is_ex {
            if payload.len() < 5 {
                return Err("enhanced RTMP payload too short".into());
            }
            match &payload[1..5] {
                b"hvc1" => Ok(StreamCodec::H265),
                b"avc1" => Ok(StreamCodec::H264),
                b"av01" => Ok(StreamCodec::AV1),
                _ => Err("unsupported FourCC".into()),
            }
        } else {
            match payload[0] & 0x0F {
                7 => Ok(StreamCodec::H264),
                _ => Err("unsupported legacy codec ID".into()),
            }
        }
    }
    async fn handshake(
        mut socket: TcpStream,
    ) -> Result<(ServerSession, TcpStream), Box<dyn Error>> {
        let mut server = Handshake::new(PeerType::Server);

        let mut c0_c1 = [0u8; 1537];
        socket.read_exact(&mut c0_c1).await?;
        let s0_s1_s2 = match server.process_bytes(&c0_c1) {
            Ok(HandshakeProcessResult::InProgress { response_bytes }) => response_bytes,
            x => return Err(format!("unexpected handshake state: {:?}", x).into()),
        };
        socket.write_all(&s0_s1_s2).await?;

        let mut c2 = [0u8; 1536];
        socket.read_exact(&mut c2).await?;
        match server.process_bytes(&c2) {
            Ok(HandshakeProcessResult::Completed { .. }) => {}
            Ok(HandshakeProcessResult::InProgress { response_bytes }) => {
                socket.write_all(&response_bytes).await?
            }
            x => return Err(format!("unexpected handshake state: {:?}", x).into()),
        }

        let (rtmp_session, init_bytes) = ServerSession::new(ServerSessionConfig::new()).unwrap();
        write_outbound(&mut socket, init_bytes).await;

        Ok((rtmp_session, socket))
    }

    pub async fn run(self) {
        let Self {
            listener,
            stream_sessions,
            db,
        } = self;
        loop {
            let connection = listener.accept().await;
            let (socket, peer_addr) = match connection {
                Ok(e) => e,
                Err(e) => {
                    tracing::error!("Error on connection accept: {}", e);
                    continue;
                }
            };
            info!(%peer_addr, "RTMP connection accepted");
            let (mut session, mut socket) = match Rtmp::handshake(socket).await {
                Ok(v) => v,
                Err(e) => {
                    warn!(%peer_addr, "RTMP handshake failed: {e}");
                    continue;
                }
            };
            info!(%peer_addr, "RTMP handshake complete");
            let (mut video_tx, mut video_rx) = broadcast::<Arc<VideoFrame>>(32);
            let (mut audio_tx, mut audio_rx) = broadcast::<Arc<OpusAudioFrame>>(32);
            // video_rx.cycle

            video_tx.set_overflow(true);
            audio_tx.set_overflow(true);
            let mut parser: Option<Box<dyn CodecParser>> = None;
            let mut stream_id: Option<i32> = None;
            let mut aac_parser = AACParser::new();
            let mut audio_proc = AudioProcesser::new();
            let db = db.clone();
            let stream_sessions = stream_sessions.clone();

            tokio::spawn(async move {
                let mut current_stream_key_id: Option<i32> = None;
                let mut codec_stamped = false;

                let cleanup = |id: i32| {
                    let stream_sessions = stream_sessions.clone();
                    let db = db.clone();
                    async move {
                        stream_sessions.remove(&id);
                        if let Ok(Some(s)) =
                            stream_session::Model::get_active_by_stream_key_id(&db, id).await
                        {
                            s.into_active_model()
                                .finish_stream_session(&db, Local::now().into())
                                .await
                                .ok();
                        }
                    }
                };

                loop {
                    let mut buf = [0u8; 4096];
                    const IDLE_TIMEOUT: Duration = Duration::from_secs(15);
                    let n = match timeout(IDLE_TIMEOUT, socket.read(&mut buf)).await {
                        Ok(Ok(0)) => {
                            debug!("RTMP connection closed by peer");
                            if let Some(id) = current_stream_key_id {
                                cleanup(id).await;
                            }
                            return;
                        }
                        Ok(Ok(n)) => n,
                        Ok(Err(e)) => {
                            warn!("RTMP read error: {e}");
                            if let Some(id) = current_stream_key_id {
                                cleanup(id).await;
                            }
                            return;
                        }
                        Err(_) => {
                            warn!("RTMP timed out (15 sec no packets) id: {:?}", stream_id);
                            if let Some(id) = current_stream_key_id {
                                cleanup(id).await;
                            }
                            return;
                        }
                    };
                    let events = match session.handle_input(&buf[..n]) {
                        Ok(e) => e,
                        Err(e) => {
                            warn!("RTMP session error: {e}");
                            if let Some(id) = current_stream_key_id {
                                cleanup(id).await;
                            }
                            return;
                        }
                    };
                    // Blankly using it, so it doesnt drop
                    video_rx.is_closed();
                    audio_rx.is_closed();

                    for event in events {
                        match event {
                            ServerSessionResult::OutboundResponse(p) => {
                                socket.write_all(&p.bytes).await.unwrap();
                            }
                            ServerSessionResult::RaisedEvent(e) => match e {
                                ServerSessionEvent::ConnectionRequested { request_id, .. } => {
                                    debug!("RTMP ConnectionRequested, accepting");
                                    let reply = session.accept_request(request_id).unwrap();
                                    write_outbound(&mut socket, reply).await;
                                }
                                ServerSessionEvent::PublishStreamRequested {
                                    request_id,
                                    stream_key,
                                    ..
                                } => {
                                    info!(stream_key = %stream_key, "publish stream requested");
                                    let key =
                                        entity::stream_key::Entity::find_by_key(&db, &stream_key)
                                            .await;

                                    let key = if let Ok(Some(key)) = key {
                                        key
                                    } else {
                                        warn!(stream_key = %stream_key, "stream key not found, rejecting");
                                        let reply = session
                                            .reject_request(request_id, "", "Stream key invalid")
                                            .unwrap();
                                        write_outbound(&mut socket, reply).await;
                                        break;
                                    };

                                    let already_live =
                                        stream_session::Model::get_active_by_stream_key_id(
                                            &db, key.id,
                                        )
                                        .await
                                        .unwrap();
                                    if already_live.is_some() {
                                        warn!(stream_key_id = key.id, label = %key.label, "stream key already live, rejecting duplicate publish");
                                        let reply = session
                                            .reject_request(
                                                request_id,
                                                "",
                                                "You're already streaming...",
                                            )
                                            .unwrap();
                                        write_outbound(&mut socket, reply).await;
                                        break;
                                    }

                                    info!(stream_key_id = key.id, label = %key.label, "stream started");
                                    current_stream_key_id = Some(key.id);
                                    stream_sessions.insert(
                                        key.id,
                                        StreamSession {
                                            stream_key_id: key.id,
                                            stream_key_label: key.label,
                                            frame_channel: video_tx.clone(),
                                            audio_channel: audio_tx.clone(),
                                            codec: None,
                                            started_at: Utc::now(),
                                        },
                                    );

                                    let reply = session.accept_request(request_id).unwrap();
                                    stream_session::Model::create_stream_session(
                                        &db,
                                        key.id,
                                        Local::now().into(),
                                    )
                                    .await
                                    .unwrap();
                                    stream_id = Some(key.id);
                                    write_outbound(&mut socket, reply).await;
                                }
                                ServerSessionEvent::PublishStreamFinished {
                                    stream_key, ..
                                } => {
                                    info!(stream_key = %stream_key, "publish stream finished");
                                    match entity::stream_key::Entity::find_by_key(&db, &stream_key)
                                        .await
                                    {
                                        Ok(Some(key)) => cleanup(key.id).await,
                                        Ok(None) => {
                                            warn!(stream_key = %stream_key, "publish finished for unknown stream key");
                                        }
                                        Err(e) => {
                                            warn!(stream_key = %stream_key, "failed to look up stream key on finish: {e}");
                                        }
                                    }
                                    current_stream_key_id = None;
                                }
                                // TODO: We can totally replace the broadcast with a
                                // circular_buff
                                // Arc<Vec<ArcSwap<Frame>>>
                                ServerSessionEvent::VideoDataReceived {
                                    data, timestamp, ..
                                } => match Self::parse_video_codec(&data) {
                                    Ok(codec) => {
                                        if !codec_stamped {
                                            if let Some(id) = current_stream_key_id {
                                                if let Some(mut session) =
                                                    stream_sessions.get_mut(&id)
                                                {
                                                    session.codec = Some(codec.clone());
                                                }
                                            }
                                            codec_stamped = true;
                                        }
                                        match codec {
                                            StreamCodec::H264 => {
                                                let p = parser.get_or_insert_with(|| {
                                                    Box::new(H264CodecParser::new())
                                                });
                                                if let Some(frame) = p.parse(&data, timestamp.value)
                                                {
                                                    video_tx.broadcast(Arc::new(frame)).await.ok();
                                                }
                                            }
                                            StreamCodec::H265 => {
                                                let p = parser.get_or_insert_with(|| {
                                                    Box::new(H265CodecParser::new())
                                                });
                                                if let Some(frame) = p.parse(&data, timestamp.value)
                                                {
                                                    video_tx.broadcast(Arc::new(frame)).await.ok();
                                                }
                                            }
                                            StreamCodec::AV1 => {
                                                let p = parser.get_or_insert_with(|| {
                                                    Box::new(Av1CodecParser::new())
                                                });
                                                let pkt_type = data[0] & 0x0F;
                                                match p.parse(&data, timestamp.value) {
                                                    Some(frame) => {
                                                        debug!(
                                                            pkt_type,
                                                            is_keyframe = frame.is_keyframe,
                                                            ts = frame.timestamp_ms,
                                                            bytes = frame.data.len(),
                                                            "AV1 frame → broadcast"
                                                        );
                                                        video_tx
                                                            .broadcast(Arc::new(frame))
                                                            .await
                                                            .ok();
                                                    }
                                                    None => {
                                                        debug!(
                                                            pkt_type,
                                                            "AV1 packet produced no frame (seq header or unknown type)"
                                                        );
                                                    }
                                                }
                                            }
                                        }
                                    }
                                    Err(_err) => {
                                        warn!("");
                                    }
                                },
                                ServerSessionEvent::AudioDataReceived {
                                    data, timestamp, ..
                                } => {
                                    // Consume the non-Send error before any await point.
                                    let opus_frames: Vec<_> =
                                        match aac_parser.parse(&data, timestamp.value) {
                                            Err(e) => {
                                                warn!("AAC parse error: {}", e);
                                                vec![]
                                            }
                                            Ok(frame) => frame
                                                .map(|f| audio_proc.encode(f))
                                                .unwrap_or_default(),
                                        };
                                    for frame in opus_frames {
                                        audio_tx.broadcast(Arc::new(frame)).await.ok();
                                    }
                                }
                                _ => {}
                            },
                            _ => {}
                        }
                    }
                }
            });
        }
    }
}