// Copyright 2026 Mirakurun contributors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.

use std::collections::{BTreeMap, HashMap, HashSet};

use mirakurun_types::{
    Program, ProgramAudio, ProgramGenre, ProgramRelatedItem, ProgramRelatedItemType, ProgramSeries,
    ProgramVideo, ProgramVideoResolution, ProgramVideoType, program_id,
};

use crate::{
    arib::decode as decode_arib,
    filter::SectionAssembler,
    ts::{Packet, PacketFramer, crc32_mpeg2},
};

const PID_EIT: u16 = 0x0012;

#[derive(Debug, Default)]
pub struct EitCollector {
    framer: PacketFramer,
    assembler: SectionAssembler,
    programs: HashMap<u64, Program>,
    states: HashMap<u64, ProgramState>,
    sections_seen: HashSet<(u8, u16, u8, u8)>,
    packet_count: u64,
    section_count: u64,
}

impl EitCollector {
    pub fn push(&mut self, chunk: &[u8]) {
        let mut packets = Vec::new();
        self.framer.push(chunk, |packet| packets.push(*packet));
        for packet_bytes in packets {
            self.packet_count = self.packet_count.saturating_add(1);
            let Ok(packet) = Packet::new(&packet_bytes) else {
                continue;
            };
            if packet.pid() != PID_EIT {
                continue;
            }
            for section in self.assembler.push(packet) {
                if parse_eit_section(&section, &mut self.programs, &mut self.states)
                    && self.sections_seen.insert((
                        section[0],
                        u16::from_be_bytes([section[3], section[4]]),
                        (section[5] >> 1) & 0x1f,
                        section[6],
                    ))
                {
                    self.section_count = self.section_count.saturating_add(1);
                }
            }
        }
    }

    #[must_use]
    pub const fn packet_count(&self) -> u64 {
        self.packet_count
    }

    #[must_use]
    pub const fn section_count(&self) -> u64 {
        self.section_count
    }

    #[must_use]
    pub fn into_programs(self) -> Vec<Program> {
        let mut programs = self.programs.into_values().collect::<Vec<_>>();
        programs.sort_by_key(|program| (program.start_at, program.id));
        programs
    }
}

#[derive(Debug, Default)]
struct VersionRecord(HashMap<u8, u8>);

impl VersionRecord {
    fn should_update(&self, table_id: u8, version: u8) -> bool {
        if self.0.contains_key(&0x4e) && table_id != 0x4e {
            return false;
        }
        if self.0.contains_key(&0x4f) && !matches!(table_id, 0x4e | 0x4f) {
            return false;
        }
        self.0.get(&table_id).copied() != Some(version)
    }

    fn mark(&mut self, table_id: u8, version: u8) {
        self.0.insert(table_id, version);
    }
}

#[derive(Debug, Default)]
struct ProgramState {
    initialized: bool,
    present: bool,
    following: bool,
    event: VersionRecord,
    short: VersionRecord,
    extended: ExtendedState,
    component: VersionRecord,
    content: VersionRecord,
    audio: HashMap<u8, VersionRecord>,
    audios: BTreeMap<u8, ProgramAudio>,
    series: VersionRecord,
    group: HashMap<u8, VersionRecord>,
    groups: BTreeMap<u8, Vec<ProgramRelatedItem>>,
}

#[derive(Debug, Default)]
struct ExtendedState {
    versions: VersionRecord,
    parts: BTreeMap<u8, Vec<ExtendedItem>>,
    last_descriptor: Option<u8>,
    done: bool,
}

#[derive(Debug)]
struct ExtendedItem {
    description: Vec<u8>,
    value: Vec<u8>,
}

fn parse_eit_section(
    section: &[u8],
    programs: &mut HashMap<u64, Program>,
    states: &mut HashMap<u64, ProgramState>,
) -> bool {
    if section.len() < 18
        || !matches!(section[0], 0x4e..=0x6f)
        || section[5] & 0x01 == 0
        || crc32_mpeg2(section) != 0
    {
        return false;
    }
    let table_id = section[0];
    let version = (section[5] >> 1) & 0x1f;
    let is_present_following = matches!(section[0], 0x4e | 0x4f);
    if is_present_following && section[6] > 1 {
        return false;
    }
    let is_present = is_present_following && section[6] == 0;
    let is_following = is_present_following && section[6] == 1;
    let service_id = u16::from_be_bytes([section[3], section[4]]);
    let network_id = u16::from_be_bytes([section[10], section[11]]);
    let Some(end) = section.len().checked_sub(4) else {
        return false;
    };
    let mut offset = 14;
    while offset + 12 <= end {
        let event_id = u16::from_be_bytes([section[offset], section[offset + 1]]);
        let start = &section[offset + 2..offset + 7];
        let duration = &section[offset + 7..offset + 10];
        let descriptor_length =
            (usize::from(section[offset + 10] & 0x0f) << 8) | usize::from(section[offset + 11]);
        let Some(descriptor_end) = offset
            .checked_add(12)
            .and_then(|value| value.checked_add(descriptor_length))
        else {
            return false;
        };
        if descriptor_end > end {
            return false;
        }
        let id = program_id(network_id, service_id, event_id);
        let start_at = decode_mjd_time(start);
        if !programs.contains_key(&id) && start_at.is_none() {
            offset = descriptor_end;
            continue;
        }
        let program = programs.entry(id).or_insert_with(|| Program {
            id,
            event_id,
            service_id,
            network_id,
            start_at: start_at.unwrap_or_default(),
            duration: decode_bcd_duration(duration).unwrap_or(1),
            is_free: section[offset + 10] & 0x10 == 0,
            name: None,
            description: None,
            genres: None,
            video: None,
            audios: None,
            series: None,
            extended: None,
            related_items: None,
        });
        let state = states.entry(id).or_default();
        let update_event = !state.initialized
            || (!state.present && is_present)
            || (!state.following && is_following)
            || state.event.should_update(table_id, version);
        if update_event {
            if let Some(start_at) = start_at {
                program.start_at = start_at;
                program.duration = decode_bcd_duration(duration).unwrap_or(1);
                program.is_free = section[offset + 10] & 0x10 == 0;
            }
            state.initialized = true;
            state.present = is_present;
            state.following = is_following;
            state.event.mark(table_id, version);
        }
        parse_descriptors(
            &section[offset + 12..descriptor_end],
            program,
            state,
            table_id,
            version,
        );
        offset = descriptor_end;
    }
    true
}

fn parse_descriptors(
    mut descriptors: &[u8],
    program: &mut Program,
    state: &mut ProgramState,
    table_id: u8,
    version: u8,
) {
    while descriptors.len() >= 2 {
        let length = usize::from(descriptors[1]);
        let Some(end) = 2usize.checked_add(length) else {
            return;
        };
        if end > descriptors.len() {
            return;
        }
        let body = &descriptors[2..end];
        match descriptors[0] {
            0x4d if state.short.should_update(table_id, version) => {
                if let Some((name, description)) = parse_short_event(body) {
                    state.short.mark(table_id, version);
                    program.name = Some(name);
                    program.description = Some(description);
                }
            }
            0x4e => parse_extended_event(body, program, &mut state.extended, table_id, version),
            0x50 if state.component.should_update(table_id, version) => {
                if let Some(video) = parse_video_component(body) {
                    state.component.mark(table_id, version);
                    program.video = Some(video);
                }
            }
            0x54 if state.content.should_update(table_id, version) => {
                if let Some(genres) = parse_content(body) {
                    state.content.mark(table_id, version);
                    program.genres = Some(genres);
                }
            }
            0xc4 => parse_audio_component(body, program, state, table_id, version),
            0xd5 if state.series.should_update(table_id, version) => {
                if let Some(series) = parse_series(body) {
                    state.series.mark(table_id, version);
                    program.series = Some(series);
                }
            }
            0xd6 => parse_event_group(body, program, state, table_id, version),
            _ => {}
        }
        descriptors = &descriptors[end..];
    }
}

fn parse_short_event(body: &[u8]) -> Option<(String, String)> {
    if body.len() < 5 {
        return None;
    }
    let name_length = usize::from(body[3]);
    let text_length_offset = 4usize.checked_add(name_length)?;
    if text_length_offset >= body.len() {
        return None;
    }
    let text_length = usize::from(body[text_length_offset]);
    let text_start = text_length_offset + 1;
    let text_end = text_start.checked_add(text_length)?;
    if text_end > body.len() {
        return None;
    }
    Some((
        decode_arib(&body[4..text_length_offset]),
        decode_arib(&body[text_start..text_end]),
    ))
}

fn parse_extended_event(
    body: &[u8],
    program: &mut Program,
    state: &mut ExtendedState,
    table_id: u8,
    version: u8,
) {
    if body.len() < 6 {
        return;
    }
    let descriptor_number = body[0] >> 4;
    let last_descriptor = body[0] & 0x0f;
    if state.versions.should_update(table_id, version) {
        state.versions.mark(table_id, version);
        state.parts.clear();
        state.last_descriptor = Some(last_descriptor);
        state.done = false;
    } else if state.done {
        return;
    }
    let items_length = usize::from(body[4]);
    let Some(items_end) = 5usize.checked_add(items_length) else {
        return;
    };
    if items_end > body.len() {
        return;
    }
    let Some(items) = parse_extended_items(&body[5..items_end]) else {
        return;
    };
    state.parts.entry(descriptor_number).or_insert(items);
    let Some(last) = state.last_descriptor else {
        return;
    };
    if !(0..=last).all(|number| state.parts.contains_key(&number)) {
        return;
    }
    program.extended = Some(build_extended(&state.parts));
    state.done = true;
}

fn parse_extended_items(mut items: &[u8]) -> Option<Vec<ExtendedItem>> {
    let mut parsed = Vec::new();
    while !items.is_empty() {
        let description_length = usize::from(items[0]);
        let description_end = 1usize.checked_add(description_length)?;
        if description_end >= items.len() {
            return None;
        }
        let item_length = usize::from(items[description_end]);
        let item_start = description_end + 1;
        let item_end = item_start.checked_add(item_length)?;
        if item_end > items.len() {
            return None;
        }
        parsed.push(ExtendedItem {
            description: items[1..description_end].to_vec(),
            value: items[item_start..item_end].to_vec(),
        });
        items = &items[item_end..];
    }
    Some(parsed)
}

fn build_extended(parts: &BTreeMap<u8, Vec<ExtendedItem>>) -> BTreeMap<String, String> {
    let mut buffers: BTreeMap<String, Vec<Vec<u8>>> = BTreeMap::new();
    let mut current_key = String::new();
    for items in parts.values() {
        for item in items {
            let decoded_key = decode_arib(&item.description);
            if !decoded_key.is_empty() {
                current_key = decoded_key;
            }
            let values = buffers.entry(current_key.clone()).or_default();
            if item.description.is_empty() && !values.is_empty() {
                if let Some(previous) = values.last_mut() {
                    previous.extend_from_slice(&item.value);
                }
            } else {
                values.push(item.value.clone());
            }
        }
    }
    buffers
        .into_iter()
        .map(|(key, values)| {
            let value = values
                .iter()
                .map(|value| decode_arib(value))
                .collect::<Vec<_>>()
                .join("\n\n");
            (key, value)
        })
        .collect()
}

fn parse_video_component(body: &[u8]) -> Option<ProgramVideo> {
    if body.len() < 2 {
        return None;
    }
    let stream_content = body[0] & 0x0f;
    let component_type = body[1];
    let video_type = match stream_content {
        1 => ProgramVideoType::Mpeg2,
        5 => ProgramVideoType::H264,
        9 => ProgramVideoType::H265,
        _ => return None,
    };
    let resolution = match component_type {
        0x01..=0x04 => ProgramVideoResolution::I480,
        0x83 => ProgramVideoResolution::P4320,
        0x91..=0x94 => ProgramVideoResolution::P2160,
        0xa1..=0xa4 => ProgramVideoResolution::P480,
        0xb1..=0xb4 => ProgramVideoResolution::I1080,
        0xc1..=0xc4 => ProgramVideoResolution::P720,
        0xd1..=0xd4 => ProgramVideoResolution::P240,
        0xe1..=0xe4 => ProgramVideoResolution::P1080,
        _ => return None,
    };
    Some(ProgramVideo {
        video_type,
        resolution,
        stream_content,
        component_type,
    })
}

fn parse_content(body: &[u8]) -> Option<Vec<ProgramGenre>> {
    let genres = body
        .chunks_exact(2)
        .map(|content| ProgramGenre {
            lv1: content[0] >> 4,
            lv2: content[0] & 0x0f,
            un1: content[1] >> 4,
            un2: content[1] & 0x0f,
        })
        .collect::<Vec<_>>();
    (!genres.is_empty()).then_some(genres)
}

fn parse_audio_component(
    body: &[u8],
    program: &mut Program,
    state: &mut ProgramState,
    table_id: u8,
    version: u8,
) {
    if body.len() < 9 {
        return;
    }
    let component_tag = body[2];
    let versions = state.audio.entry(component_tag).or_default();
    if !versions.should_update(table_id, version) {
        return;
    }
    let multilingual = body[5] & 0x80 != 0;
    if multilingual && body.len() < 12 {
        return;
    }
    let mut langs = vec![language_code(&body[6..9])];
    if multilingual {
        langs.push(language_code(&body[9..12]));
    }
    let sampling_rate = match (body[5] >> 1) & 0x07 {
        1 => 16_000,
        2 => 22_050,
        3 => 24_000,
        5 => 32_000,
        6 => 44_100,
        7 => 48_000,
        _ => return,
    };
    versions.mark(table_id, version);
    state.audios.insert(
        component_tag,
        ProgramAudio {
            component_type: body[1],
            component_tag,
            is_main: body[5] & 0x40 != 0,
            sampling_rate,
            langs,
        },
    );
    program.audios = Some(state.audios.values().cloned().collect());
}

fn language_code(value: &[u8]) -> String {
    match value {
        b"jpn" | b"eng" | b"deu" | b"fra" | b"ita" | b"rus" | b"zho" | b"kor" | b"spa" => {
            String::from_utf8_lossy(value).into_owned()
        }
        _ => "etc".into(),
    }
}

fn parse_series(body: &[u8]) -> Option<ProgramSeries> {
    if body.len() < 8 {
        return None;
    }
    let expires_at = if body[2] & 0x01 != 0 {
        decode_mjd_date(&body[3..5]).unwrap_or(-1)
    } else {
        -1
    };
    let episodes = (u32::from(body[5]) << 16) | (u32::from(body[6]) << 8) | u32::from(body[7]);
    Some(ProgramSeries {
        id: u16::from_be_bytes([body[0], body[1]]),
        repeat: body[2] >> 4,
        pattern: (body[2] >> 1) & 0x07,
        expires_at,
        episode: u16::try_from(episodes >> 12).ok()?,
        last_episode: u16::try_from(episodes & 0x0fff).ok()?,
        name: decode_arib(&body[8..]),
    })
}

fn parse_event_group(
    body: &[u8],
    program: &mut Program,
    state: &mut ProgramState,
    table_id: u8,
    version: u8,
) {
    let Some(&header) = body.first() else {
        return;
    };
    let group_type = header >> 4;
    let event_count = usize::from(header & 0x0f);
    let versions = state.group.entry(group_type).or_default();
    if !versions.should_update(table_id, version) {
        return;
    }
    let local_end = 1usize.saturating_add(event_count.saturating_mul(4));
    if local_end > body.len() {
        return;
    }
    let item_type = match group_type {
        1 => ProgramRelatedItemType::Shared,
        2 | 4 => ProgramRelatedItemType::Relay,
        _ => ProgramRelatedItemType::Movement,
    };
    let mut items = Vec::new();
    if matches!(group_type, 4 | 5) {
        for event in body[local_end..].chunks_exact(8) {
            items.push(ProgramRelatedItem {
                item_type,
                network_id: Some(u16::from_be_bytes([event[0], event[1]])),
                service_id: u16::from_be_bytes([event[4], event[5]]),
                event_id: u16::from_be_bytes([event[6], event[7]]),
            });
        }
    } else {
        for event in body[1..local_end].chunks_exact(4) {
            items.push(ProgramRelatedItem {
                item_type,
                network_id: None,
                service_id: u16::from_be_bytes([event[0], event[1]]),
                event_id: u16::from_be_bytes([event[2], event[3]]),
            });
        }
    }
    versions.mark(table_id, version);
    state.groups.insert(group_type, items);
    program.related_items = Some(state.groups.values().flatten().cloned().collect());
}

fn decode_mjd_time(value: &[u8]) -> Option<u64> {
    if value.len() != 5 || value.iter().all(|byte| *byte == 0xff) {
        return None;
    }
    let mjd = i64::from(u16::from_be_bytes([value[0], value[1]]));
    let hour = i64::from(decode_bcd(value[2])?);
    let minute = i64::from(decode_bcd(value[3])?);
    let second = i64::from(decode_bcd(value[4])?);
    if hour > 23 || minute > 59 || second > 59 {
        return None;
    }
    let unix_seconds = (mjd - 40_587)
        .checked_mul(86_400)?
        .checked_add((hour - 9).checked_mul(3600)?)?
        .checked_add(minute.checked_mul(60)?)?
        .checked_add(second)?;
    u64::try_from(unix_seconds.checked_mul(1000)?).ok()
}

fn decode_mjd_date(value: &[u8]) -> Option<i64> {
    if value.len() != 2 || value.iter().all(|byte| *byte == 0xff) {
        return None;
    }
    let mjd = i64::from(u16::from_be_bytes([value[0], value[1]]));
    (mjd - 40_587)
        .checked_mul(86_400)?
        .checked_sub(9 * 3_600)?
        .checked_mul(1000)
}

fn decode_bcd_duration(value: &[u8]) -> Option<u64> {
    if value.len() != 3 || value.iter().all(|byte| *byte == 0xff) {
        return None;
    }
    let hours = u64::from(decode_bcd(value[0])?);
    let minutes = u64::from(decode_bcd(value[1])?);
    let seconds = u64::from(decode_bcd(value[2])?);
    if minutes > 59 || seconds > 59 {
        return None;
    }
    hours
        .checked_mul(3600)?
        .checked_add(minutes.checked_mul(60)?)?
        .checked_add(seconds)?
        .checked_mul(1000)
}

fn decode_bcd(value: u8) -> Option<u8> {
    let high = value >> 4;
    let low = value & 0x0f;
    (high <= 9 && low <= 9).then_some(high * 10 + low)
}

#[cfg(test)]
mod tests {
    use crate::ts::{PACKET_SIZE, crc32_mpeg2};
    use mirakurun_types::ProgramRelatedItemType;

    use super::{EitCollector, VersionRecord, decode_arib, decode_bcd_duration, decode_mjd_time};

    #[test]
    fn decodes_common_arib_character_sets() {
        assert_eq!(decode_arib(&[0x0e, b'T', b'e', b's', b't']), "Test");
        assert_eq!(decode_arib(&[0x19, 0x22]), "あ");
        assert_eq!(decode_arib(&[0x1d, 0x22]), "ア");
    }

    #[test]
    fn converts_arib_time_fields() {
        assert_eq!(decode_mjd_time(&[0x9e, 0x8b, 0x09, 0x00, 0x00]), Some(0));
        assert_eq!(decode_bcd_duration(&[0x01, 0x30, 0x00]), Some(5_400_000));
    }

    #[test]
    fn collects_program_from_eit_packet() {
        let short_event = [
            0x4d, 0x0d, b'j', b'p', b'n', 0x05, 0x0e, b'T', b'e', b's', b't', 0x03, 0x0e, b'O',
            b'K',
        ];
        let mut section = vec![
            0x4e, 0xb0, 0x00, 0x00, 0x65, 0xc1, 0x00, 0x00, 0x00, 0x01, 0x7f, 0xf0, 0x00, 0x4e,
            0x00, 0x01, 0x9e, 0x8b, 0x09, 0x00, 0x00, 0x01, 0x00, 0x00, 0x80, 0x0f,
        ];
        section.extend_from_slice(&short_event);
        let section_length = section.len() - 3 + 4;
        section[1] = 0xb0 | u8::try_from(section_length >> 8).expect("section length");
        section[2] = u8::try_from(section_length & 0xff).expect("section length");
        let crc = crc32_mpeg2(&section);
        section.extend_from_slice(&crc.to_be_bytes());

        let mut packet = [0xff; PACKET_SIZE];
        packet[0] = 0x47;
        packet[1] = 0x40;
        packet[2] = 0x12;
        packet[3] = 0x10;
        packet[4] = 0;
        packet[5..5 + section.len()].copy_from_slice(&section);

        let mut collector = EitCollector::default();
        collector.push(&packet);
        let programs = collector.into_programs();
        assert_eq!(programs.len(), 1);
        assert_eq!(programs[0].name.as_deref(), Some("Test"));
        assert_eq!(programs[0].description.as_deref(), Some("OK"));
        assert_eq!(programs[0].duration, 3_600_000);
        assert!(programs[0].is_free);
    }

    #[test]
    fn parses_audio_series_and_related_program_descriptors() {
        let mut descriptors = vec![
            0xc4, 0x09, 0x02, 0x03, 0x10, 0x0f, 0xff, 0x4e, b'j', b'p', b'n',
        ];
        let series_name = [0x0e, b'S', b'e', b'r', b'i', b'e', b's'];
        descriptors.extend_from_slice(&[
            0xd5,
            u8::try_from(8 + series_name.len()).unwrap_or(0),
            0x12,
            0x34,
            0x14,
            0xff,
            0xff,
            0x00,
            0x50,
            0x0c,
        ]);
        descriptors.extend_from_slice(&series_name);
        descriptors.extend_from_slice(&[0xd6, 0x05, 0x11, 0x00, 0x65, 0x00, 0x02]);

        let packet = eit_packet(0x4e, 0, &descriptors);
        let mut collector = EitCollector::default();
        collector.push(&packet);
        let programs = collector.into_programs();
        let program = &programs[0];
        let audio = &program.audios.as_ref().expect("audio descriptor")[0];
        assert_eq!(audio.component_type, 0x03);
        assert_eq!(audio.component_tag, 0x10);
        assert!(audio.is_main);
        assert_eq!(audio.sampling_rate, 48_000);
        assert_eq!(audio.langs, ["jpn"]);
        let series = program.series.as_ref().expect("series descriptor");
        assert_eq!(series.id, 0x1234);
        assert_eq!(series.repeat, 1);
        assert_eq!(series.pattern, 2);
        assert_eq!(series.episode, 5);
        assert_eq!(series.last_episode, 12);
        assert_eq!(series.name, "Series");
        let related = &program.related_items.as_ref().expect("event group")[0];
        assert_eq!(related.item_type, ProgramRelatedItemType::Shared);
        assert_eq!(related.network_id, None);
        assert_eq!(related.service_id, 101);
        assert_eq!(related.event_id, 2);
    }

    #[test]
    fn present_following_versions_take_priority_over_schedule_tables() {
        let mut versions = VersionRecord::default();
        assert!(versions.should_update(0x50, 1));
        versions.mark(0x50, 1);
        assert!(versions.should_update(0x4e, 1));
        versions.mark(0x4e, 1);
        assert!(!versions.should_update(0x50, 2));
        assert!(versions.should_update(0x4e, 2));
    }

    fn eit_packet(table_id: u8, version: u8, descriptors: &[u8]) -> [u8; PACKET_SIZE] {
        let descriptor_length = descriptors.len();
        let mut section = vec![
            table_id,
            0xb0,
            0x00,
            0x00,
            0x65,
            0xc1 | ((version & 0x1f) << 1),
            0x00,
            0x00,
            0x00,
            0x01,
            0x7f,
            0xf0,
            0x00,
            table_id,
            0x00,
            0x01,
            0x9e,
            0x8b,
            0x09,
            0x00,
            0x00,
            0x01,
            0x00,
            0x00,
            0x80 | u8::try_from(descriptor_length >> 8).unwrap_or(0),
            u8::try_from(descriptor_length & 0xff).unwrap_or(0),
        ];
        section.extend_from_slice(descriptors);
        let section_length = section.len() - 3 + 4;
        section[1] = 0xb0 | u8::try_from(section_length >> 8).unwrap_or(0);
        section[2] = u8::try_from(section_length & 0xff).unwrap_or(0);
        let crc = crc32_mpeg2(&section);
        section.extend_from_slice(&crc.to_be_bytes());

        let mut packet = [0xff; PACKET_SIZE];
        packet[0] = 0x47;
        packet[1] = 0x40;
        packet[2] = 0x12;
        packet[3] = 0x10;
        packet[4] = 0;
        packet[5..5 + section.len()].copy_from_slice(&section);
        packet
    }
}