use argparse::{ArgumentParser, Store, StoreFalse, StoreTrue};
use colored::*;
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};
use tracing::{debug, info, Level};

use spvirit_scry::common::stats::IfStats;
use spvirit_scry::utils::packet_processor_common::{create_capture, parse_packet, PvaReassembler};
use spvirit_codec::epics_decode::{PvaCommands, PvaPacket, PvaPacketCommand};
use spvirit_codec::spvirit_state::{ConnectionKey, PvaStateConfig, PvaStateTracker};

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let mut verbose = true;
    let mut device = String::from("bond10gbe");
    let mut debug_mode = false;
    let mut pv_name = String::new();
    let mut server_ip = String::new();
    let mut limit = 50usize;
    let mut duration_secs = 20u64;
    let mut dump_file = String::new();
    let mut dump_hex = false;
    let mut dump_raw = String::new();
    let mut dump_validation = false;

    {
        let mut ap = ArgumentParser::new();
        ap.set_description("Capture a successful PVA GET for a specific PV");
        ap.refer(&mut verbose)
            .add_option(&["-q", "--quiet"], StoreFalse, "Be quiet");
        ap.refer(&mut device)
            .add_option(&["-i", "--device"], Store, "Device to capture from");
        ap.refer(&mut debug_mode)
            .add_option(&["-d", "--debug"], StoreTrue, "Show debug messages");
        ap.refer(&mut pv_name)
            .add_option(&["--pv"], Store, "PV name to track (required)");
        ap.refer(&mut server_ip)
            .add_option(&["--server-ip"], Store, "Optional server IP filter");
        ap.refer(&mut limit)
            .add_option(&["--limit"], Store, "Max messages to print (default 50)");
        ap.refer(&mut duration_secs).add_option(
            &["--duration"],
            Store,
            "Max capture time in seconds (default 20)",
        );
        ap.refer(&mut dump_file).add_option(
            &["--dump-file"],
            Store,
            "Write hex dumps to this file",
        );
        ap.refer(&mut dump_hex).add_option(
            &["--dump-hex"],
            StoreTrue,
            "Hex dump PVA frames for matched PV",
        );
        ap.refer(&mut dump_raw).add_option(
            &["--dump-raw"],
            Store,
            "Write raw frames (len+bytes) to this file",
        );
        ap.refer(&mut dump_validation).add_option(
            &["--dump-validation"],
            StoreTrue,
            "Dump validation frames regardless of PV match",
        );
        ap.parse_args_or_exit();
    }

    if pv_name.trim().is_empty() {
        eprintln!("--pv is required");
        std::process::exit(1);
    }

    if verbose {
        if debug_mode {
            tracing_subscriber::fmt::fmt()
                .with_max_level(Level::DEBUG)
                .init();
        } else {
            tracing_subscriber::fmt::fmt()
                .with_max_level(Level::INFO)
                .init();
        }
        info!("Starting PVA GET capture");
        info!("{}", format!("Using device {}", device).green());
        if !server_ip.is_empty() {
            info!("Filtering server IP: {}", server_ip);
        }
        if dump_hex {
            info!("Hex dump enabled");
        }
        if !dump_file.is_empty() {
            info!("Dump file: {}", dump_file);
        }
        if !dump_raw.is_empty() {
            info!("Raw dump file: {}", dump_raw);
        }
        if dump_validation {
            info!("Dump validation enabled");
        }
    }

    let if_stats = Arc::new(Mutex::new(IfStats::new()));
    let pva_cmds = PvaCommands::new();
    let mut pva_state = PvaStateTracker::new(PvaStateConfig::new(40_000, 300));
    let mut cap = create_capture(&device);
    let _ = cap.filter("udp port 5076 or tcp port 5075", true);
    let mut reassembler = PvaReassembler::new();

    let start = Instant::now();
    let mut printed = 0usize;
    let mut saw_get_data = false;
    let mut dump_writer: Option<std::fs::File> = if !dump_file.is_empty() {
        Some(
            std::fs::OpenOptions::new()
                .create(true)
                .append(true)
                .open(&dump_file)?,
        )
    } else {
        None
    };
    let mut raw_writer: Option<std::fs::File> = if !dump_raw.is_empty() {
        Some(
            std::fs::OpenOptions::new()
                .create(true)
                .append(true)
                .open(&dump_raw)?,
        )
    } else {
        None
    };

    'outer: loop {
        if start.elapsed() > Duration::from_secs(duration_secs) {
            info!("Capture timeout reached");
            break;
        }
        if printed >= limit {
            info!("Message limit reached");
            break;
        }

        let packet = match cap.next_packet() {
            Ok(p) => p,
            Err(_) => continue,
        };
        let parsed = match parse_packet(&packet.data) {
            Some(parsed) => parsed,
            None => continue,
        };
        if !server_ip.is_empty()
            && parsed.source != server_ip
            && parsed.destination != server_ip
        {
            continue;
        }
        if parsed.payload_length == 0 {
            continue;
        }

        let payload = &packet.data[parsed.payload_offset
            ..parsed.payload_offset + parsed.payload_length];
        let messages = reassembler.process(&parsed, payload);

        for msg_data in messages {
        let pva_message_slice = &msg_data;
        let message_command = pva_cmds.get_command(msg_data[3]);
        if message_command == "Unknown" {
            continue;
        }

        let mut packet_pva = PvaPacket::new(pva_message_slice);
        let is_server = packet_pva.header.flags.is_server;
        let mut decoded = packet_pva.decode_payload();

        let conn_key = ConnectionKey::from_parts(
            &parsed.source,
            parsed.src_port,
            &parsed.destination,
            parsed.dst_port,
        );

        if let Some(ref mut cmd) = decoded {
            if let Some(ref ck) = conn_key {
                match cmd {
                    PvaPacketCommand::CreateChannel(payload) => {
                        if is_server {
                            pva_state.on_create_channel_response(ck, payload.cid, payload.sid);
                        } else {
                            for (cid, pv) in &payload.channels {
                                pva_state.on_create_channel_request(ck, *cid, pv.clone());
                            }
                        }
                    }
                    PvaPacketCommand::Search(ref search_payload) => {
                        // Cache CID → PV name from SEARCH messages
                        if !search_payload.pv_requests.is_empty() {
                            let source_ip: Option<std::net::IpAddr> =
                                parsed.source.parse().ok();
                            pva_state.on_search(
                                &search_payload.pv_requests,
                                source_ip,
                            );
                        }
                    }
                    PvaPacketCommand::DestroyChannel(payload) => {
                        pva_state.on_destroy_channel(ck, payload.cid, payload.sid);
                    }
                    PvaPacketCommand::Op(op_payload) => {
                        if !is_server && (op_payload.subcmd & 0x08) != 0 {
                            pva_state.on_op_init_request(
                                ck,
                                op_payload.sid_or_cid,
                                op_payload.ioid,
                                op_payload.command,
                            );
                        } else if is_server && (op_payload.subcmd & 0x08) != 0 {
                            if let Some(ref intro) = op_payload.introspection {
                                pva_state.on_op_init_response(
                                    ck,
                                    op_payload.ioid,
                                    Some(intro.clone()),
                                );
                            }
                        }
                    }
                    _ => {}
                }
            }
        }

        if dump_validation {
            if let Some(PvaPacketCommand::ConnectionValidation(_)) = decoded {
                printed += 1;
                let dir = if is_server { "S->C" } else { "C->S" };
                info!(
                    "{} VALIDATION {}:{} -> {}:{}",
                    dir,
                    parsed.source,
                    parsed.src_port,
                    parsed.destination,
                    parsed.dst_port
                );
                if dump_hex {
                    let summary = format!(
                        "{} VALIDATION cmd={} flags=0x{:02x} payload_len={}",
                        dir,
                        packet_pva.header.command,
                        packet_pva.header.flags.raw,
                        packet_pva.header.payload_length
                    );
                    let hex = hex_dump(pva_message_slice, 16);
                    let dump = format!("{}\n{}\n", summary, hex);
                    print!("{}", dump);
                    if let Some(ref mut f) = dump_writer {
                        use std::io::Write;
                        let _ = f.write_all(dump.as_bytes());
                    }
                }
            }
            if let Some(PvaPacketCommand::ConnectionValidated(_)) = decoded {
                printed += 1;
                let dir = if is_server { "S->C" } else { "C->S" };
                info!(
                    "{} VALIDATED {}:{} -> {}:{}",
                    dir,
                    parsed.source,
                    parsed.src_port,
                    parsed.destination,
                    parsed.dst_port
                );
                if dump_hex {
                    let summary = format!(
                        "{} VALIDATED cmd={} flags=0x{:02x} payload_len={}",
                        dir,
                        packet_pva.header.command,
                        packet_pva.header.flags.raw,
                        packet_pva.header.payload_length
                    );
                    let hex = hex_dump(pva_message_slice, 16);
                    let dump = format!("{}\n{}\n", summary, hex);
                    print!("{}", dump);
                    if let Some(ref mut f) = dump_writer {
                        use std::io::Write;
                        let _ = f.write_all(dump.as_bytes());
                    }
                }
            }
        }

        let pv_context = match decoded {
            Some(PvaPacketCommand::Op(ref op_payload)) => {
                if let Some(ref ck) = conn_key {
                    pva_state.resolve_pv_name(ck, op_payload.sid_or_cid, op_payload.ioid)
                } else {
                    None
                }
            }
            Some(PvaPacketCommand::CreateChannel(ref payload)) => {
                if !payload.channels.is_empty() {
                    Some(payload.channels[0].1.clone())
                } else {
                    None
                }
            }
            _ => None,
        };

        if let Some(pv) = pv_context {
            if pv == pv_name {
                printed += 1;
                let dir = if is_server { "S->C" } else { "C->S" };
                info!(
                    "{} {} {}:{} -> {}:{} pv={}",
                    dir,
                    message_command,
                    parsed.source,
                    parsed.src_port,
                    parsed.destination,
                    parsed.dst_port,
                    pv
                );

                if dump_hex {
                    let summary = format!(
                        "{} {} cmd={} flags=0x{:02x} payload_len={} pv={}",
                        dir,
                        message_command,
                        packet_pva.header.command,
                        packet_pva.header.flags.raw,
                        packet_pva.header.payload_length,
                        pv
                    );
                    let hex = hex_dump(pva_message_slice, 16);
                    let dump = format!("{}\n{}\n", summary, hex);
                    print!("{}", dump);
                    if let Some(ref mut f) = dump_writer {
                        use std::io::Write;
                        let _ = f.write_all(dump.as_bytes());
                    }
                }
                if let Some(ref mut f) = raw_writer {
                    use std::io::Write;
                    let len = (pva_message_slice.len() as u32).to_le_bytes();
                    let _ = f.write_all(&len);
                    let _ = f.write_all(pva_message_slice);
                }

                if let Some(PvaPacketCommand::Op(op_payload)) = decoded {
                    if is_server && op_payload.command == 10 && op_payload.subcmd == 0x00 {
                        saw_get_data = true;
                    }
                }

                if saw_get_data {
                    info!("Captured GET data response for {}", pv_name);
                    break 'outer;
                }
            }
        }
        } // for msg_data in messages
    }

    Ok(())
}

fn hex_dump(bytes: &[u8], width: usize) -> String {
    let mut out = String::new();
    let mut offset = 0usize;
    while offset < bytes.len() {
        let end = std::cmp::min(offset + width, bytes.len());
        out.push_str(&format!("{:04x} ", offset));
        for i in offset..end {
            out.push_str(&format!("{:02x} ", bytes[i]));
        }
        out.push('\n');
        offset = end;
    }
    out
}
