//! DXGI Desktop Duplication 采集（Windows 8+）。
//! 帧处理：BGRA → openh264 H.264 编码（屏幕内容实时档，5Mbps）→ 线格式。
//! 支持主控端 KeyFrameReq 强制关键帧（解码异常恢复）。

use anyhow::{anyhow, Result};
use common::{encode_video_frame, VideoFrameHeader};
use openh264::encoder::{
    BitRate, Encoder, EncoderConfig, FrameRate, IntraFramePeriod, RateControlMode, UsageType,
};
use openh264::formats::{BgraSliceU8, YUVBuffer};
use openh264::OpenH264API;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
use tokio::sync::watch::Sender;
use windows::core::Interface;
use windows::Win32::Foundation::HMODULE;
use windows::Win32::Graphics::Direct3D::{
    D3D_DRIVER_TYPE_HARDWARE, D3D_FEATURE_LEVEL_11_0,
};
use windows::Win32::Graphics::Direct3D11::{
    D3D11CreateDevice, ID3D11Device, ID3D11DeviceContext, ID3D11Resource, ID3D11Texture2D,
    D3D11_CREATE_DEVICE_FLAG, D3D11_CPU_ACCESS_READ, D3D11_MAP, D3D11_MAP_READ,
    D3D11_MAPPED_SUBRESOURCE, D3D11_SDK_VERSION, D3D11_TEXTURE2D_DESC, D3D11_USAGE_STAGING,
};
use windows::Win32::Graphics::Dxgi::Common::{DXGI_FORMAT_B8G8R8A8_UNORM, DXGI_SAMPLE_DESC};
use windows::Win32::Graphics::Dxgi::{
    IDXGIDevice, IDXGIOutput1, IDXGIOutputDuplication, IDXGIResource, DXGI_ERROR_WAIT_TIMEOUT,
    DXGI_OUTDUPL_DESC, DXGI_OUTDUPL_FRAME_INFO,
};

use std::time::{Duration, Instant};

pub struct Duplication {
    _device: ID3D11Device,
    context: ID3D11DeviceContext,
    staging_res: ID3D11Resource,
    dup: IDXGIOutputDuplication,
    pub width: u32,
    pub height: u32,
}

impl Duplication {
    pub fn new() -> Result<Self> {
        unsafe {
            let mut device: Option<ID3D11Device> = None;
            let mut context: Option<ID3D11DeviceContext> = None;
            D3D11CreateDevice(
                None,
                D3D_DRIVER_TYPE_HARDWARE,
                HMODULE::default(),
                D3D11_CREATE_DEVICE_FLAG(0),
                Some(&[D3D_FEATURE_LEVEL_11_0]),
                D3D11_SDK_VERSION,
                Some(&mut device),
                None,
                Some(&mut context),
            )
            .map_err(|e| anyhow!("D3D11CreateDevice failed: {e}"))?;
            let device = device.ok_or_else(|| anyhow!("no d3d device"))?;
            let context = context.ok_or_else(|| anyhow!("no d3d context"))?;

            let dxgi: IDXGIDevice = device.cast()?;
            let adapter = dxgi.GetAdapter()?;
            let output = adapter
                .EnumOutputs(0)
                .map_err(|e| anyhow!("no display output: {e}"))?;
            let output1: IDXGIOutput1 = output.cast()?;
            let dup: IDXGIOutputDuplication = output1
                .DuplicateOutput(&device)
                .map_err(|e| anyhow!("DuplicateOutput failed: {e}"))?;

            let desc: DXGI_OUTDUPL_DESC = dup.GetDesc();
            let (width, height) = (desc.ModeDesc.Width, desc.ModeDesc.Height);
            anyhow::ensure!(width > 0 && height > 0, "bad desktop size");

            let sdesc = D3D11_TEXTURE2D_DESC {
                Width: width,
                Height: height,
                MipLevels: 1,
                ArraySize: 1,
                Format: DXGI_FORMAT_B8G8R8A8_UNORM,
                SampleDesc: DXGI_SAMPLE_DESC {
                    Count: 1,
                    Quality: 0,
                },
                Usage: D3D11_USAGE_STAGING,
                BindFlags: 0,
                CPUAccessFlags: D3D11_CPU_ACCESS_READ.0 as u32,
                MiscFlags: 0,
            };
            let mut staging_opt: Option<ID3D11Texture2D> = None;
            device
                .CreateTexture2D(&sdesc, None, Some(&mut staging_opt))
                .map_err(|e| anyhow!("CreateTexture2D failed: {e}"))?;
            let staging = staging_opt.ok_or_else(|| anyhow!("no staging texture"))?;
            let staging_res: ID3D11Resource = staging.cast()?;

            Ok(Self {
                _device: device,
                context,
                staging_res,
                dup,
                width,
                height,
            })
        }
    }

    /// true=拿到新帧并已复制到 staging；false=超时无更新
    pub fn acquire(&mut self, timeout_ms: u32) -> Result<bool> {
        unsafe {
            let mut info = DXGI_OUTDUPL_FRAME_INFO::default();
            let mut resource: Option<IDXGIResource> = None;
            match self.dup.AcquireNextFrame(timeout_ms, &mut info, &mut resource) {
                Ok(()) => {
                    let res = resource.ok_or_else(|| anyhow!("no desktop resource"))?;
                    let src: ID3D11Resource = res.cast()?;
                    self.context.CopyResource(&self.staging_res, &src);
                    let _ = self.dup.ReleaseFrame();
                    Ok(true)
                }
                Err(e) if e.code() == DXGI_ERROR_WAIT_TIMEOUT => Ok(false),
                Err(e) => Err(e.into()),
            }
        }
    }

    /// 把 staging 纹理映射到内存并去掉行间距
    pub fn copy_frame(&mut self) -> Result<Vec<u8>> {
        unsafe {
            let mut mapped = D3D11_MAPPED_SUBRESOURCE::default();
            self.context
                .Map(&self.staging_res, 0, D3D11_MAP_READ, 0, Some(&mut mapped))
                .map_err(|e| anyhow!("Map staging failed: {e}"))?;
            let (w, h) = (self.width as usize, self.height as usize);
            let pitch = mapped.RowPitch as usize;
            let src = std::slice::from_raw_parts(mapped.pData as *const u8, pitch * h);
            let mut out = Vec::with_capacity(w * h * 4);
            for row in 0..h {
                out.extend_from_slice(&src[row * pitch..row * pitch + w * 4]);
            }
            self.context.Unmap(&self.staging_res, 0);
            Ok(out)
        }
    }

    /// 在映射窗口内执行回调（避免 8MB 中间拷贝）。回调参数 = (像素数据, 行距)。
    /// 映射内存仅在回调返回前有效。
    pub fn with_mapped<R>(&mut self, f: impl FnOnce(&[u8], usize) -> R) -> Result<R> {
        unsafe {
            let mut mapped = D3D11_MAPPED_SUBRESOURCE::default();
            self.context
                .Map(&self.staging_res, 0, D3D11_MAP_READ, 0, Some(&mut mapped))
                .map_err(|e| anyhow!("Map staging failed: {e}"))?;
            let (w, h) = (self.width as usize, self.height as usize);
            let pitch = mapped.RowPitch as usize;
            let src = std::slice::from_raw_parts(mapped.pData as *const u8, pitch * h);
            let r = f(src, pitch);
            self.context.Unmap(&self.staging_res, 0);
            Ok(r)
        }
    }
}

#[allow(dead_code)]
pub fn desktop_size() -> Result<(u32, u32)> {
    let d = Duplication::new()?;
    Ok((d.width, d.height))
}

/// 采集主循环：独立线程运行；连接断开（通道关闭）时退出。
/// 通道为 watch 单槽（latest-wins）：发送即覆盖，网络慢时不产生积压。
pub fn run(tx: Sender<Option<Vec<u8>>>, keyreq: Arc<AtomicBool>) {
    let start = Instant::now();
    loop {
        let mut dup = match Duplication::new() {
            Ok(d) => d,
            Err(e) => {
                eprintln!("[capture] init failed: {e}, retry in 2s");
                std::thread::sleep(Duration::from_secs(2));
                continue;
            }
        };
        eprintln!(
            "[capture] duplication started {}x{}",
            dup.width, dup.height
        );

        // H.264 编码器：屏幕内容实时档 + 码率控制 + 4 线程（单线程 1080p 仅 ~20fps，
        // 多线程把编码压进 33ms 预算实现稳定 30fps）
        let mut encoder = match Encoder::with_api_config(
            OpenH264API::from_source(),
            EncoderConfig::new()
                .bitrate(BitRate::from_bps(5_000_000))
                .max_frame_rate(FrameRate::from_hz(60.0))
                .rate_control_mode(RateControlMode::Bitrate)
                .usage_type(UsageType::ScreenContentRealTime)
                .intra_frame_period(IntraFramePeriod::from_num_frames(300))
                .num_threads(4),
        ) {
            Ok(e) => e,
            Err(e) => {
                eprintln!("[capture] encoder init failed: {e}");
                return;
            }
        };

        let mut last_frame: Option<Vec<u8>> = None;
        let mut last_sent = Instant::now() - Duration::from_secs(10);
        // 编码实测 ~5ms/帧（1080p 屏幕内容实时档），60fps 只需 16.7ms 预算，
        // 够用 → 帧率上限提到 60（对标 UU 60fps；带宽由 RC 压在 5Mbps 内）
        let min_frame = Duration::from_millis(16); // 60fps 上限
        let mut scratch: Vec<u8> = Vec::new(); // 行距不连续时的复用缓冲
        // 编码耗时统计（每 120 帧输出一次均值，ms）
        let mut stat_frames = 0u64;
        let mut stat_convert_ms = 0u64;
        let mut stat_encode_ms = 0u64;
        loop {
            let iter_start = Instant::now();
            match dup.acquire(80) {
                Ok(true) => {
                    // 主控端请求关键帧（解码异常恢复）
                    let force_key = keyreq.swap(false, Ordering::Relaxed);
                    if force_key {
                        encoder.force_intra_frame();
                    }
                    // 映射窗口内直接做 BGRA→YUV（行距连续时零拷贝，否则复用 scratch）
                    let conv_start = Instant::now();
                    let dims = (dup.width as usize, dup.height as usize);
                    let conv = dup.with_mapped(|src, pitch| {
                        if pitch == dims.0 * 4 {
                            YUVBuffer::from_bgra8_source(BgraSliceU8::new(src, dims))
                        } else {
                            scratch.resize(dims.0 * dims.1 * 4, 0);
                            for (row, dst) in scratch.chunks_exact_mut(dims.0 * 4).enumerate() {
                                dst.copy_from_slice(
                                    &src[row * pitch..row * pitch + dims.0 * 4],
                                );
                            }
                            YUVBuffer::from_bgra8_source(BgraSliceU8::new(&scratch, dims))
                        }
                    });
                    match conv {
                        Ok(yuv) => {
                            stat_convert_ms += conv_start.elapsed().as_millis() as u64;
                            let enc_start = Instant::now();
                            let r = encoder.encode(&yuv);
                            stat_encode_ms += enc_start.elapsed().as_millis() as u64;
                            match r {
                                Ok(ebs) => {
                                    let payload = ebs.to_vec();
                                    let header = VideoFrameHeader {
                                        width: dup.width,
                                        height: dup.height,
                                        ts_ms: start.elapsed().as_millis() as u64,
                                    };
                                    let frame = encode_video_frame(&header, &payload);
                                    if tx.send(Some(frame.clone())).is_err() {
                                        return; // 会话结束
                                    }
                                    last_frame = Some(frame);
                                    last_sent = Instant::now();
                                    stat_frames += 1;
                                    if stat_frames % 120 == 0 {
                                        eprintln!(
                                            "[capture] avg convert {}ms, encode {}ms/frame (last 120)",
                                            stat_convert_ms / 120,
                                            stat_encode_ms / 120
                                        );
                                        stat_convert_ms = 0;
                                        stat_encode_ms = 0;
                                    }
                                }
                                Err(e) => {
                                    eprintln!("[capture] encode failed: {e}");
                                }
                            }
                        }
                        Err(e) => {
                            eprintln!("[capture] frame convert failed: {e}");
                            break;
                        }
                    }
                    // 帧间隔节流到 ~30fps
                    let el = iter_start.elapsed();
                    if el < min_frame {
                        std::thread::sleep(min_frame - el);
                    }
                }
                Ok(false) => {
                    // 桌面静止：低频重发最后一帧（重复 P 帧，解码端画面不变）
                    if last_sent.elapsed() > Duration::from_millis(500) {
                        if let Some(f) = last_frame.as_ref() {
                            if tx.send(Some(f.clone())).is_err() {
                                return;
                            }
                            last_sent = Instant::now();
                        }
                    } else {
                        std::thread::sleep(Duration::from_millis(8));
                    }
                }
                Err(e) => {
                    eprintln!("[capture] acquire failed: {e}");
                    break; // ACCESS_LOST 等，重建
                }
            }
        }
        std::thread::sleep(Duration::from_millis(300));
    }
}