use super::*;
use crate::proto::{
    CDemoStringTables, CSvcMsgCreateStringTable, CSvcMsgUpdateStringTable, Message,
};
use crate::string_table::{
    rewrite_create_string_table, rewrite_demo_string_table_items, rewrite_update_string_table,
    PackedStringTableFormat, PackedStringTableState, StringTableEntryUpdate,
};

impl<'a, R, W> DemoWriter<'a, R, W>
where
    R: BitsReader + MessageReader,
    W: Write + Seek,
{
    pub(crate) fn rewrite_string_tables(
        &mut self,
        tick: u32,
        mut msg: CDemoStringTables,
    ) -> Result<Option<(CDemoStringTables, bool)>, ParserError> {
        let mut changed = false;
        if self.rewrites_entity_fields() {
            changed |= self.rewrite_instance_baselines(&mut msg)?;
        }
        changed |= self.rewrite_demo_string_table_entries(tick, &mut msg)?;
        for rewriter in self.rewriters.iter_mut().filter(|rewriter| {
            rewriter
                .interests()
                .contains(RewriteInterests::DEMO_STRING_TABLES)
        }) {
            let ctx = &self.parser.context;
            match rewriter.rewrite_demo_string_tables(ctx, tick, &mut msg)? {
                MessageRewrite::Drop => return Ok(None),
                MessageRewrite::Keep => {}
                MessageRewrite::Rewrite => changed = true,
                MessageRewrite::Replace(bytes) => {
                    msg = CDemoStringTables::decode(bytes.as_slice())?;
                    changed = true;
                }
            }
        }
        Ok(Some((msg, changed)))
    }

    pub(crate) fn rewrite_create_string_table_entries(
        &mut self,
        tick: u32,
        msg: &mut CSvcMsgCreateStringTable,
    ) -> Result<bool, ParserError> {
        let table_name = msg.name().to_string();
        let rewrite_baselines =
            table_name == INSTANCE_BASELINE_TABLE && self.rewrites_entity_fields();
        if !self.rewrites_string_table_entries() && !rewrite_baselines {
            return Ok(false);
        }

        let table_id = self.string_table_rewrite_states.len();
        let mut state =
            PackedStringTableState::new(PackedStringTableFormat::from_create_message(msg));

        let changed = rewrite_create_string_table(msg, &mut state, |entry| {
            if rewrite_baselines {
                self.rewrite_instance_baseline_entry_update(entry, None)?;
            }
            self.rewrite_string_table_entry_with_rewriters(tick, &table_name, entry)?;
            Ok(())
        });

        let changed = changed?;

        self.ensure_string_table_rewrite_state(table_id);
        self.string_table_rewrite_states[table_id] = Some(state);
        Ok(changed)
    }

    pub(crate) fn rewrite_update_string_table_entries(
        &mut self,
        tick: u32,
        msg: &mut CSvcMsgUpdateStringTable,
    ) -> Result<bool, ParserError> {
        let table_id = msg.table_id() as usize;
        let Some(table) = self.parser.context.string_tables.tables.get(table_id) else {
            return Ok(false);
        };
        let table_name = table.name().to_string();
        let rewrite_baselines =
            table_name == INSTANCE_BASELINE_TABLE && self.rewrites_entity_fields();
        if !self.rewrites_string_table_entries() && !rewrite_baselines {
            return Ok(false);
        }
        let existing_keys = if rewrite_baselines {
            table
                .iter()
                .map(|row| row.key().to_string())
                .collect::<Vec<_>>()
        } else {
            Vec::new()
        };
        let state_from_context = PackedStringTableState::from_table(table);

        self.ensure_string_table_rewrite_state(table_id);
        if self.string_table_rewrite_states[table_id].is_none() {
            self.string_table_rewrite_states[table_id] = Some(state_from_context);
        }

        let Some(mut state) = self.string_table_rewrite_states[table_id].take() else {
            return Err(ParserError::MissingStringTableRewriteState { table_id });
        };

        let changed = rewrite_update_string_table(msg, &mut state, |entry| {
            if rewrite_baselines {
                let fallback_key = usize::try_from(entry.index())
                    .ok()
                    .and_then(|index| existing_keys.get(index))
                    .map(String::as_str);
                self.rewrite_instance_baseline_entry_update(entry, fallback_key)?;
            }
            self.rewrite_string_table_entry_with_rewriters(tick, &table_name, entry)?;
            Ok(())
        });

        self.string_table_rewrite_states[table_id] = Some(state);
        changed
    }

    pub(crate) fn rewrite_demo_string_table_entries(
        &mut self,
        tick: u32,
        msg: &mut CDemoStringTables,
    ) -> Result<bool, ParserError> {
        if !self.rewrites_string_table_entries() {
            return Ok(false);
        }

        let mut changed = false;
        for table in msg.tables.iter_mut() {
            let table_name = table.table_name().to_string();
            changed |= rewrite_demo_string_table_items(&mut table.items, |entry| {
                self.rewrite_string_table_entry_with_rewriters(tick, &table_name, entry)
            })?;
            changed |= rewrite_demo_string_table_items(&mut table.items_clientside, |entry| {
                self.rewrite_string_table_entry_with_rewriters(tick, &table_name, entry)
            })?;
        }
        Ok(changed)
    }

    fn ensure_string_table_rewrite_state(&mut self, table_id: usize) {
        if self.string_table_rewrite_states.len() <= table_id {
            self.string_table_rewrite_states
                .resize_with(table_id + 1, || None);
        }
    }

    fn rewrite_instance_baseline_entry_update(
        &mut self,
        entry: &mut StringTableEntryUpdate,
        fallback_key: Option<&str>,
    ) -> Result<(), ParserError> {
        let Some(class_id) = entry
            .key()
            .or(fallback_key)
            .and_then(|value| value.parse::<i32>().ok())
        else {
            return Ok(());
        };
        if class_id < 0 {
            return Ok(());
        }

        let Some(data) = entry.value() else {
            return Ok(());
        };
        if let Some(rewritten) = self.rewrite_instance_baseline_data(class_id, data)? {
            entry.set_value(rewritten);
        }
        Ok(())
    }

    fn rewrite_string_table_entry_with_rewriters(
        &mut self,
        tick: u32,
        table_name: &str,
        entry: &mut StringTableEntryUpdate,
    ) -> Result<(), ParserError> {
        for rewriter in self.rewriters.iter_mut().filter(|rewriter| {
            rewriter
                .interests()
                .contains(RewriteInterests::STRING_TABLE_ENTRIES)
        }) {
            let ctx = &self.parser.context;
            rewriter.rewrite_string_table_entry(ctx, tick, table_name, entry)?;
        }
        Ok(())
    }
}
