use anyhow::{bail, Result};
use clap::Subcommand;
use dfir_windows_types::{Guid, Sid};
use sddl::{AccessMaskFlag, Contains};
use serde::Serialize;

use hashbrown::{hash_map::Entry, HashMap, HashSet};
use std::str::FromStr;
use strum::Display;

use crate::{
    cli::OutputFormat,
    ntds::{DataTableRecord, InterestingAcl},
    win32_types::Rdn,
    CDatabase,
};

#[derive(Subcommand, Clone, Display, Eq, PartialEq)]
pub enum Command {
    /// list all available extended rights
    List {
        /// Output format
        #[clap(value_enum, short('F'), long("format"), default_value_t = OutputFormat::Csv)]
        format: OutputFormat,
    },

    /// list all subjects (and resp. objects) with some specific extended right
    Subjects {
        /// name or GUID of the extended right to search for, or nothing if
        /// you want to search for objects
        extended_right: Option<String>,

        /// SID of the object that you are interested in. If you specify the SID,
        /// only subjects that have the requested extended right will be shown.
        #[clap(long("object"))]
        object: Option<Sid>,
    },
}

impl Command {
    pub fn run<'info, 'db>(&self, database: &CDatabase<'info, 'db>) -> Result<()> {
        match self {
            Command::List { format } => self.list_extended_rights(database, format),
            Command::Subjects {
                extended_right,
                object,
            } => self.list_extended_right_subjects(
                database,
                extended_right.as_ref().map(|x| x.as_str()),
                object.as_ref(),
            ),
        }
    }

    fn list_extended_rights<'info, 'db>(
        &self,
        database: &CDatabase<'info, 'db>,
        format: &OutputFormat,
    ) -> Result<()> {
        #[derive(Serialize)]
        struct ExtendedRight {
            object_name: String,
            rights_guid: Guid,
            object_guid: Guid,
        }

        if let Ok(extended_rights) = database
            .data_table()
            .object_tree()
            .get_special_records()?
            .extended_rights_container()
            .children()
            .lock()
        {
            let mut writer = format.create_writer()?;
            for right in extended_rights.iter() {
                let object = database
                    .entry(crate::EntryId::Id(*right.record_ptr().ds_record_id()))?
                    .expect("illegal record id in object tree");

                writer.write(ExtendedRight {
                    object_name: object.att_object_name()?.to_string(),
                    rights_guid: object
                        .att_rights_guid_opt()?
                        .expect("extended right without rights GUID"),
                    object_guid: object.att_object_guid()?,
                })?;
            }
            drop(writer);
        } else {
            bail!("unable to acquire lock to extended rights container")
        }

        Ok(())
    }

    fn list_extended_right_subjects<'info, 'db>(
        &self,
        database: &CDatabase<'info, 'db>,
        extended_right: Option<&str>,
        object: Option<&Sid>,
    ) -> std::result::Result<(), anyhow::Error> {
        let format = OutputFormat::Json;

        #[derive(Serialize, Eq, PartialEq, Clone)]
        struct Subject {
            subject_name: Option<Rdn>,
            subject_guid: Option<Guid>,
            subject_sid: Sid,
            permission_for_objects: HashMap<String, HashSet<ObjectInfo>>,
        }

        #[derive(Serialize, Eq, PartialEq, Hash, Clone)]
        struct ObjectInfo {
            object_name: Option<Rdn>,
            object_guid: Option<Guid>,
            object_sid: Option<Sid>,
        }

        let right_guid = match extended_right {
            None => None,
            Some(r) => Some(self.resolve_extended_right(database, r)?),
        };

        let mut extended_rights_mapping: HashMap<Sid, Subject> = HashMap::new();

        let sd_map = self.find_sd_with_extended_right(database, right_guid)?;

        for (sd_id, acl) in sd_map.iter() {
            let mut permission_for_objects = HashSet::new();
            for obj in self.find_objects_with_sd(database, *sd_id)? {
                let object_name = obj.att_object_name_opt()?;
                let object_guid = obj.att_object_guid_opt()?;
                let object_sid = obj.att_object_sid_opt()?;

                if let Some(object) = object {
                    if let Some(object_sid) = object_sid {
                        if object == &object_sid {
                            let object_sid = Some(object_sid);
                            let obj_info = ObjectInfo {
                                object_name,
                                object_guid,
                                object_sid,
                            };
                            permission_for_objects.insert(obj_info);
                        }
                    }
                } else {
                    let obj_info = ObjectInfo {
                        object_name,
                        object_guid,
                        object_sid,
                    };
                    permission_for_objects.insert(obj_info);
                }
            }

            if !permission_for_objects.is_empty() {
                for sid in acl.sids().iter().cloned() {
                    let entry = database
                        .data_table()
                        .data_table()
                        .metadata()
                        .entry_with_sid(&sid);
                    match extended_rights_mapping.entry(sid.clone()) {
                        Entry::Occupied(mut sid_entry) => {
                            for object in permission_for_objects.clone().into_iter() {
                                match sid_entry
                                    .get_mut()
                                    .permission_for_objects
                                    .entry(extended_right.unwrap_or("<ALL_RIGHTS>").to_owned())
                                {
                                    Entry::Occupied(mut entry) => {
                                        (*entry.get_mut()).insert(object);
                                    }
                                    Entry::Vacant(entry) => {
                                        let mut permission_for_objects = HashSet::new();
                                        permission_for_objects.insert(object);
                                        entry.insert(permission_for_objects);
                                    }
                                };
                            }
                        }
                        Entry::Vacant(sid_entry) => {
                            let mut pfs = HashMap::new();
                            pfs.insert(
                                extended_right.unwrap_or("<ALL_RIGHTS>").to_owned(),
                                permission_for_objects.clone(),
                            );
                            let subject = Subject {
                                subject_sid: sid,
                                subject_name: entry.map(|e| e.rdn()).cloned(),
                                subject_guid: entry.and_then(|e| e.guid().as_ref()).cloned(),
                                permission_for_objects: pfs,
                            };
                            sid_entry.insert(subject);
                        }
                    }
                }
            }
        }

        let mut writer = format.create_writer()?;
        writer.write(extended_rights_mapping)?;
        drop(writer);

        Ok(())
    }

    fn resolve_extended_right<'info, 'db>(
        &self,
        database: &CDatabase<'info, 'db>,
        extended_right: &str,
    ) -> std::result::Result<Guid, anyhow::Error> {
        if let Ok(guid) = Guid::from_str(extended_right) {
            Ok(guid)
        } else if let Ok(extended_rights) = database
            .data_table()
            .object_tree()
            .get_special_records()?
            .extended_rights_container()
            .children()
            .lock()
        {
            for right in extended_rights.iter() {
                let object = database
                    .entry(crate::EntryId::Id(*right.record_ptr().ds_record_id()))?
                    .expect("illegal record id in object tree");
                if object.att_object_name()?.as_str() == extended_right {
                    return Ok(object
                        .att_rights_guid_opt()?
                        .expect("extended right without rights GUID"));
                }
            }
            bail!("unknown extended right: '{extended_right}'");
        } else {
            bail!("unable to acquire lock to Extended Rights container")
        }
    }

    /// If `right_guid` is `None`, then search for ACEs with the
    /// CONTROL_ACCESS flag set, but without rights guid (which means
    /// that all extended rights are permitted)
    fn find_sd_with_extended_right<'info, 'db>(
        &self,
        database: &CDatabase<'info, 'db>,
        right_guid: Option<Guid>,
    ) -> std::result::Result<HashMap<i64, InterestingAcl>, anyhow::Error> {
        Ok(database.data_table().find_acls(move |ace| {
            if ace.header().mask().contains(AccessMaskFlag::CONTROL_ACCESS) {
                match right_guid {
                    None if ace.object_type().is_none() => Some(ace.sid().sid()),
                    Some(guid) => {
                        if let Some(object_type) = ace.object_type() {
                            if guid == *object_type {
                                Some(ace.sid().sid())
                            } else {
                                None
                            }
                        } else {
                            None
                        }
                    }
                    _ => None,
                }
            } else {
                None
            }
        }))
    }

    fn find_objects_with_sd<'info, 'db>(
        &self,
        database: &'info CDatabase<'info, 'db>,
        sd_id: i64,
    ) -> anyhow::Result<Vec<DataTableRecord<'info, 'db>>> {
        let mut objects = Vec::new();
        if let Some(record_ids) = database
            .data_table()
            .data_table()
            .metadata()
            .sd_objects()
            .get(&sd_id)
        {
            for record_id in record_ids {
                let record = database
                    .data_table()
                    .entry(crate::EntryId::Id(*record_id))?
                    .expect("illegal record pointer");
                objects.push(record);
            }
        }
        Ok(objects)
    }
}
