use std::cell::RefCell;
use std::collections::{BTreeMap, btree_map::Entry as BTreeMapEntry};
use std::io::{Read, Write};
use std::num::NonZeroUsize;
use std::rc::Rc;

use serde::ser::{Serialize, SerializeMap, SerializeSeq, Serializer};
use serde_json::{Value, json};
use simd_csv::{ByteRecord, StringRecord};
use simd_json::{
    prelude::{TypedValue, ValueAsScalar, ValueType},
    tape::Value as TapeValue,
};

use crate::moonblade::{Path, Step};
use crate::select::Selection;

struct SerializableTapeValue<'tape, 'input>(TapeValue<'tape, 'input>);

impl Serialize for SerializableTapeValue<'_, '_> {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        match self.0.value_type() {
            ValueType::Null => serializer.serialize_unit(),
            ValueType::Bool => serializer.serialize_bool(self.0.as_bool().unwrap()),
            ValueType::String => serializer.serialize_str(self.0.as_str().unwrap()),
            ValueType::I64 => serializer.serialize_i64(self.0.as_i64().unwrap()),
            ValueType::U64 => serializer.serialize_u64(self.0.as_u64().unwrap()),
            ValueType::F64 => serializer.serialize_f64(self.0.as_f64().unwrap()),
            ValueType::Array => {
                let array = self.0.as_array().unwrap();

                let mut seq = serializer.serialize_seq(Some(array.len()))?;

                for value in array.iter() {
                    seq.serialize_element(&Self(value))?;
                }

                seq.end()
            }
            ValueType::Object => {
                let object = self.0.as_object().unwrap();

                let mut map = serializer.serialize_map(Some(object.len()))?;

                for (k, v) in object.iter() {
                    map.serialize_entry(k, &Self(v))?;
                }

                map.end()
            }
            _ => unreachable!(),
        }
    }
}

pub trait GetPathOwned {
    fn get_path_owned(self, path: &Path) -> Option<Self>
    where
        Self: Sized;
}

impl GetPathOwned for Value {
    fn get_path_owned(self, path: &Path) -> Option<Self> {
        let mut value = self;

        for step in path.iter() {
            match step {
                &Step::Index(mut i) => {
                    let array = value.as_array_mut()?;

                    if i < 0 {
                        i += array.len() as isize;
                    }

                    if i < 0 || i >= array.len() as isize {
                        return None;
                    } else {
                        value = array.swap_remove(i as usize);
                    }
                }
                Step::Key(key) => {
                    let object = value.as_object_mut()?;
                    value = object.remove(key)?;
                }
            }
        }

        Some(value)
    }
}

impl GetPathOwned for TapeValue<'_, '_> {
    fn get_path_owned(self, path: &Path) -> Option<Self> {
        let mut value = self;

        for step in path.iter() {
            match step {
                &Step::Index(mut i) => {
                    if let Some(array) = value.as_array() {
                        if i < 0 {
                            i += array.len() as isize;
                        }

                        value = array.get(i as usize)?;
                    } else {
                        return None;
                    }
                }
                Step::Key(key) => {
                    if let Some(object) = value.as_object() {
                        value = object.get(key.as_str())?;
                    } else {
                        return None;
                    }
                }
            }
        }

        Some(value)
    }
}

// pub trait GetPathBorrowed {
//     fn get_path_borrowed<'s>(&'s self, path: &Path) -> Option<&'s Self>
//     where
//         Self: Sized;
// }

// impl GetPathBorrowed for BorrowedValue<'_> {
//     fn get_path_borrowed<'s>(&'s self, path: &Path) -> Option<&'s Self> {
//         let mut value = self;

//         for step in path.iter() {
//             match step {
//                 &Step::Index(mut i) => {
//                     if let BorrowedValue::Array(array) = value {
//                         if i < 0 {
//                             i += array.len() as isize;
//                         }

//                         value = array.get(i as usize)?;
//                     } else {
//                         return None;
//                     }
//                 }
//                 Step::Key(key) => {
//                     if let BorrowedValue::Object(object) = value {
//                         value = object.get(key.as_str())?;
//                     } else {
//                         return None;
//                     }
//                 }
//             }
//         }

//         Some(value)
//     }
// }

#[derive(Default)]
pub struct AttributeNameInterner {
    strings: Vec<Rc<String>>,
    map: BTreeMap<Rc<String>, usize>,
}

impl AttributeNameInterner {
    pub fn register(&mut self, name: String) -> usize {
        use BTreeMapEntry::*;

        let name = Rc::new(name);
        let next_id = self.strings.len();

        match self.map.entry(name.clone()) {
            Occupied(entry) => *entry.get(),
            Vacant(entry) => {
                self.strings.push(name.clone());
                entry.insert(next_id);
                next_id
            }
        }
    }

    fn get(&self, id: usize) -> &str {
        &self.strings[id]
    }
}

thread_local! {
    pub static INTERNER: RefCell<AttributeNameInterner> = RefCell::new(AttributeNameInterner::default());
}

#[derive(Debug, Default)]
pub struct Attributes {
    entries: Vec<(usize, Value)>,
}

impl Attributes {
    pub fn is_empty(&self) -> bool {
        self.entries.is_empty()
    }

    pub fn with_capacity(capacity: usize) -> Self {
        Self {
            entries: Vec::with_capacity(capacity),
        }
    }

    pub fn insert(&mut self, key: &str, value: Value) {
        let attr_name_id = INTERNER.with_borrow_mut(|interner| interner.register(key.to_string()));
        self.entries.push((attr_name_id, value));
    }

    pub fn iter(&self) -> impl Iterator<Item = &(usize, Value)> {
        self.entries.iter()
    }

    pub fn get(&self, id: usize) -> Option<&Value> {
        self.entries
            .iter()
            .find_map(|entry| if entry.0 == id { Some(&entry.1) } else { None })
    }
}

impl Serialize for Attributes {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        let mut map = serializer.serialize_map(Some(self.entries.len()))?;

        for (k, v) in self.entries.iter() {
            INTERNER.with_borrow(|interner| map.serialize_entry(interner.get(*k), v))?;
        }

        map.end()
    }
}

#[derive(Debug, Default, Clone)]
pub struct OmittableAttributes {
    entries: Vec<(usize, Option<Value>)>,
}

impl OmittableAttributes {
    pub fn from_headers<'a>(headers: impl Iterator<Item = &'a str>) -> Self {
        let mut entries = Vec::new();

        for h in headers {
            INTERNER
                .with_borrow_mut(|interner| entries.push((interner.register(h.to_string()), None)));
        }

        Self { entries }
    }

    pub fn values_mut(&mut self) -> impl Iterator<Item = &mut Option<Value>> {
        self.entries.iter_mut().map(|(_, v)| v)
    }
}

impl Serialize for OmittableAttributes {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        let mut map = serializer.serialize_map(Some(self.entries.len()))?;

        for (k, opt) in self.entries.iter() {
            if let Some(v) = opt {
                INTERNER.with_borrow(|interner| map.serialize_entry(interner.get(*k), v))?;
            }
        }

        map.end()
    }
}

// NOTE: we keep a depth on `Delve` and `Pop` to be able to skip absent keys efficiently
#[derive(Debug, Clone)]
enum JSONTraversalState {
    Delve(String, usize),
    Emit,
    Pop(usize),
}

type JSONTraversalStack = Vec<JSONTraversalState>;

fn traverse_to_build_stack(
    value: &Value,
    stack: &mut JSONTraversalStack,
    reorder_keys: bool,
    depth: usize,
) {
    match value {
        Value::Object(map) => {
            let mut items = map.iter().collect::<Vec<_>>();

            // NOTE: we put scalar values first, then nested ones and we also sort by key
            if reorder_keys {
                items.sort_by_key(|i| {
                    (
                        if matches!(i.1, Value::Object(_)) {
                            1
                        } else {
                            0
                        },
                        i.0,
                    )
                });
            }

            for (k, v) in items {
                stack.push(JSONTraversalState::Delve(k.to_string(), depth));

                traverse_to_build_stack(v, stack, reorder_keys, depth + 1);

                stack.push(JSONTraversalState::Pop(depth));
            }
        }
        _ => stack.push(JSONTraversalState::Emit),
    };
}

fn headers_from_stack(stack: &JSONTraversalStack) -> ByteRecord {
    let mut record = ByteRecord::new();

    // Single scalar early return
    if stack.len() == 1 {
        record.push_field(b"value");
        return record;
    }

    let mut path: Vec<&str> = Vec::new();

    for state in stack {
        match state {
            JSONTraversalState::Delve(key, _) => {
                path.push(key.as_str());
            }
            JSONTraversalState::Emit => {
                record.push_field(path.join(".").as_bytes());
            }
            JSONTraversalState::Pop(_) => {
                path.pop();
            }
        }
    }

    record
}

fn traverse_value_with_stack<F>(mut value: &Value, stack: &JSONTraversalStack, mut callback: F)
where
    F: FnMut(&Value),
{
    let mut working_stack: Vec<&Value> = vec![];

    let mut i: usize = 0;

    'outer: while i < stack.len() {
        let state = &stack[i];

        match state {
            JSONTraversalState::Delve(key, depth) => {
                working_stack.push(value);

                match value.as_object().and_then(|o| o.get(key)) {
                    None => {
                        i += 1;

                        while i < stack.len() {
                            let next_state = &stack[i];

                            match next_state {
                                JSONTraversalState::Emit => callback(&Value::Null),
                                JSONTraversalState::Pop(target_depth) if target_depth == depth => {
                                    continue 'outer;
                                }
                                _ => (),
                            };

                            i += 1;
                        }
                    }
                    Some(next_value) => {
                        value = next_value;
                    }
                }
            }
            JSONTraversalState::Pop(_) => {
                value = working_stack.pop().expect("cannot pop");
            }
            JSONTraversalState::Emit => {
                callback(value);
            }
        }

        i += 1;
    }
}

// fn traverse_borrowed_value_with_stack<F>(
//     mut value: &BorrowedValue,
//     stack: &JSONTraversalStack,
//     mut callback: F,
// ) where
//     F: FnMut(&BorrowedValue),
// {
//     let mut working_stack: Vec<&BorrowedValue> = vec![];

//     let mut i: usize = 0;

//     'outer: while i < stack.len() {
//         let state = &stack[i];

//         match state {
//             JSONTraversalState::Delve(key, depth) => {
//                 working_stack.push(value);

//                 let as_object = match value {
//                     BorrowedValue::Object(object) => Some(object),
//                     _ => None,
//                 };

//                 match as_object.and_then(|o| o.get(key.as_str())) {
//                     None => {
//                         i += 1;

//                         while i < stack.len() {
//                             let next_state = &stack[i];

//                             match next_state {
//                                 JSONTraversalState::Emit => {
//                                     callback(&BorrowedValue::Static(simd_json::StaticNode::Null))
//                                 }
//                                 JSONTraversalState::Pop(target_depth) if target_depth == depth => {
//                                     continue 'outer;
//                                 }
//                                 _ => (),
//                             };

//                             i += 1;
//                         }
//                     }
//                     Some(next_value) => {
//                         value = next_value;
//                     }
//                 }
//             }
//             JSONTraversalState::Pop(_) => {
//                 value = working_stack.pop().expect("cannot pop");
//             }
//             JSONTraversalState::Emit => {
//                 callback(value);
//             }
//         }

//         i += 1;
//     }
// }

fn traverse_tape_value_with_stack<F>(
    mut value: TapeValue,
    stack: &JSONTraversalStack,
    mut callback: F,
) where
    F: FnMut(TapeValue),
{
    let mut working_stack: Vec<TapeValue> = vec![];

    let mut i: usize = 0;

    'outer: while i < stack.len() {
        let state = &stack[i];

        match state {
            JSONTraversalState::Delve(key, depth) => {
                working_stack.push(value);

                match value.get(key.as_str()) {
                    None => {
                        i += 1;

                        while i < stack.len() {
                            let next_state = &stack[i];

                            match next_state {
                                JSONTraversalState::Emit => callback(TapeValue::null()),
                                JSONTraversalState::Pop(target_depth) if target_depth == depth => {
                                    continue 'outer;
                                }
                                _ => (),
                            };

                            i += 1;
                        }
                    }
                    Some(next_value) => {
                        value = next_value;
                    }
                }
            }
            JSONTraversalState::Pop(_) => {
                value = working_stack.pop().expect("cannot pop");
            }
            JSONTraversalState::Emit => {
                callback(value);
            }
        }

        i += 1;
    }
}

#[inline]
fn fill_record_from_value(value: &Value, record: &mut ByteRecord, stack: &JSONTraversalStack) {
    record.clear();

    traverse_value_with_stack(value, stack, |v| match v {
        Value::Null => record.push_field(b""),
        Value::Bool(b) => record.push_field(if *b { b"true" } else { b"false" }),
        Value::String(s) => record.push_field(s.as_bytes()),
        Value::Number(n) => record.fmt_field(n),
        Value::Array(l) => record.write_field(|view| serde_json::to_writer(view, l).unwrap()),
        Value::Object(o) => record.write_field(|view| serde_json::to_writer(view, o).unwrap()),
    });
}

// #[inline]
// fn fill_record_from_borrowed_value(
//     value: &BorrowedValue,
//     record: &mut ByteRecord,
//     stack: &JSONTraversalStack,
// ) {
//     record.clear();

//     traverse_borrowed_value_with_stack(value, stack, |v| match v {
//         BorrowedValue::Static(simd_json::StaticNode::Null) => record.push_field(b""),
//         BorrowedValue::Static(simd_json::StaticNode::Bool(b)) => {
//             record.push_field(if *b { b"true" } else { b"false" });
//         }
//         BorrowedValue::String(s) => record.push_field(s.as_bytes()),
//         BorrowedValue::Static(simd_json::StaticNode::F64(f)) => {
//             record.fmt_field(f);
//         }
//         BorrowedValue::Static(simd_json::StaticNode::U64(i)) => {
//             record.fmt_field(i);
//         }
//         BorrowedValue::Static(simd_json::StaticNode::I64(i)) => {
//             record.fmt_field(i);
//         }
//         BorrowedValue::Array(l) => {
//             record.write_field(|view| serde_json::to_writer(view, l).unwrap())
//         }
//         BorrowedValue::Object(o) => {
//             record.write_field(|view| serde_json::to_writer(view, o).unwrap())
//         }
//     });
// }

#[inline]
fn fill_record_from_tape_value(
    value: TapeValue,
    record: &mut ByteRecord,
    stack: &JSONTraversalStack,
) {
    record.clear();

    // TODO: we could work on the Tape itself without bothering with clunky TapeValue
    traverse_tape_value_with_stack(value, stack, |v| match v.value_type() {
        ValueType::Null => {
            record.push_field(b"");
        }
        ValueType::Bool => {
            record.push_field(if v.as_bool().unwrap() {
                b"true"
            } else {
                b"false"
            });
        }
        ValueType::String => {
            record.push_field(v.as_str().unwrap().as_bytes());
        }
        ValueType::F64 => {
            record.fmt_field(&v.as_f64().unwrap());
        }
        ValueType::I64 => {
            record.fmt_field(&v.as_i64().unwrap());
        }
        ValueType::U64 => {
            record.fmt_field(&v.as_u64().unwrap());
        }
        ValueType::Array | ValueType::Object => record
            .write_field(|view| serde_json::to_writer(view, &SerializableTapeValue(v)).unwrap()),
        _ => unreachable!(),
    });
}

fn merge_with_owned(a: &mut Value, b: &Value) {
    if let Value::Object(a) = a {
        if let Value::Object(b) = b {
            for (k, v) in b {
                merge_with_owned(a.entry(k).or_insert(Value::Null), v);
            }

            return;
        }
    }

    *a = b.clone();
}

pub struct JSONTabularizer<W: Write> {
    writer: simd_csv::Writer<W>,
    harmonized_value: Value,
    sample_size: Option<usize>,
    sample: Vec<Value>,
    flushed: bool,
    output_record: ByteRecord,
    stack: JSONTraversalStack,
    reorder_keys: bool,
}

impl<W: Write> JSONTabularizer<W> {
    pub fn from_writer(writer: simd_csv::Writer<W>, sample_size: Option<NonZeroUsize>) -> Self {
        Self {
            writer,
            harmonized_value: Value::Null,
            sample_size: sample_size.map(|limit| limit.get()),
            sample: match sample_size {
                Some(limit) => Vec::with_capacity(limit.get()),
                None => Vec::new(),
            },
            flushed: true,
            output_record: ByteRecord::new(),
            stack: JSONTraversalStack::new(),
            reorder_keys: false,
        }
    }

    pub fn set_model(&mut self, model: &Value) {
        self.sample_size = Some(0);
        self.flushed = false;

        merge_with_owned(&mut self.harmonized_value, model);
    }

    pub fn reorder_keys(&mut self) {
        self.reorder_keys = true;
    }

    pub fn flush(&mut self) -> simd_csv::Result<()> {
        if self.flushed {
            return Ok(());
        }

        traverse_to_build_stack(
            &self.harmonized_value,
            &mut self.stack,
            self.reorder_keys,
            0,
        );

        self.writer
            .write_byte_record(&headers_from_stack(&self.stack))?;

        for value in self.sample.iter() {
            fill_record_from_value(value, &mut self.output_record, &self.stack);
            self.writer.write_byte_record(&self.output_record)?;
        }

        self.writer.flush()?;

        self.flushed = true;

        Ok(())
    }

    pub fn is_sampling(&self) -> bool {
        self.sample_size.is_none() || self.sample.len() < self.sample_size.unwrap()
    }

    pub fn process(&mut self, value: Value) -> simd_csv::Result<()> {
        // Sampling
        if self.is_sampling() {
            self.flushed = false;
            merge_with_owned(&mut self.harmonized_value, &value);
            self.sample.push(value);
            return Ok(());
        }

        self.flush()?;

        fill_record_from_value(&value, &mut self.output_record, &self.stack);
        self.writer.write_byte_record(&self.output_record)?;

        Ok(())
    }

    // pub fn process_borrowed_no_sampling(&mut self, value: &BorrowedValue) -> simd_csv::Result<()> {
    //     self.flush()?;

    //     fill_record_from_borrowed_value(value, &mut self.output_record, &self.stack);
    //     self.writer.write_byte_record(&self.output_record)?;

    //     Ok(())
    // }

    pub fn process_tape_no_sampling(&mut self, value: TapeValue) -> simd_csv::Result<()> {
        self.flush()?;

        fill_record_from_tape_value(value, &mut self.output_record, &self.stack);
        self.writer.write_byte_record(&self.output_record)?;

        Ok(())
    }
}

const JSON_MAX_SAFE_INTEGER: i64 = 9007199254740991;

#[derive(Debug, Clone, Copy)]
pub enum JSONEmptyMode {
    Null,
    Empty,
    Omit,
}

#[derive(Debug, Clone, Copy, PartialEq)]
pub enum JSONType {
    Null,
    String,
    Integer,
    Float,
}

impl JSONType {
    fn merge(self, other: Self) -> Self {
        match self {
            Self::Null => other,
            Self::String => self,
            Self::Integer => match other {
                Self::Float | Self::String => other,
                Self::Integer | Self::Null => self,
            },
            Self::Float => match other {
                Self::Float | Self::Integer | Self::Null => self,
                Self::String => other,
            },
        }
    }
}

#[derive(Debug)]
struct JSONTypeInferrence {
    json_types: Vec<JSONType>,
    empty_mode: JSONEmptyMode,
}

impl JSONTypeInferrence {
    pub fn new(columns: usize, empty_mode: JSONEmptyMode) -> Self {
        let mut json_types = Vec::with_capacity(columns);

        for _ in 0..columns {
            json_types.push(JSONType::Null);
        }

        Self {
            json_types,
            empty_mode,
        }
    }

    fn set_string(&mut self, index: usize) {
        self.json_types[index] = JSONType::String;
    }

    fn infer(value: &str) -> JSONType {
        if value.is_empty() {
            return JSONType::Null;
        }

        if let Ok(integer) = value.parse::<i64>() {
            if integer.abs() <= JSON_MAX_SAFE_INTEGER {
                return JSONType::Integer;
            } else {
                return JSONType::String;
            }
        }

        if value.parse::<f64>().is_ok() {
            return JSONType::Float;
        }

        JSONType::String
    }

    fn cast(&self, value: &str, json_type: JSONType) -> Option<Value> {
        if value.is_empty() {
            return match self.empty_mode {
                JSONEmptyMode::Omit => None,
                JSONEmptyMode::Null => Some(Value::Null),
                JSONEmptyMode::Empty => Some(json!("")),
            };
        }

        match json_type {
            JSONType::String | JSONType::Null => Some(json!(value)),
            JSONType::Float => {
                if let Ok(float) = value.parse::<f64>() {
                    Some(json!(float))
                } else {
                    Some(json!(value))
                }
            }
            JSONType::Integer => {
                if let Ok(integer) = value.parse::<i64>() {
                    if integer.abs() <= JSON_MAX_SAFE_INTEGER {
                        return Some(json!(integer));
                    }
                }

                Some(json!(value))
            }
        }
    }

    fn process<'a>(&mut self, values: impl Iterator<Item = &'a str>) {
        for (json_type, value) in self.json_types.iter_mut().zip(values) {
            let new_json_type = Self::infer(value);
            *json_type = json_type.merge(new_json_type);
        }
    }
}

#[derive(Debug)]
pub struct JSONTypeInferrenceBuffer {
    inferrence: JSONTypeInferrence,
    buffer: Vec<StringRecord>,
    capacity: Option<usize>,
    selection: Selection,
}

impl JSONTypeInferrenceBuffer {
    pub fn new(
        selection: Selection,
        buffer_size: Option<usize>,
        empty_mode: JSONEmptyMode,
    ) -> Self {
        Self {
            inferrence: JSONTypeInferrence::new(selection.len(), empty_mode),
            buffer: match buffer_size {
                Some(capacity) => Vec::with_capacity(capacity),
                None => Vec::new(),
            },
            capacity: buffer_size,
            selection,
        }
    }

    pub fn with_columns(
        columns: usize,
        buffer_size: Option<usize>,
        empty_mode: JSONEmptyMode,
    ) -> Self {
        Self::new(Selection::full(columns), buffer_size, empty_mode)
    }

    pub fn set_string(&mut self, index: usize) {
        self.inferrence.set_string(index);
    }

    pub fn read<R: Read>(
        &mut self,
        reader: &mut simd_csv::Reader<R>,
    ) -> Result<(), simd_csv::Error> {
        match self.capacity {
            Some(capacity) => {
                for result in reader.records().take(capacity) {
                    self.process(result?);
                }
            }
            None => {
                for result in reader.records() {
                    self.process(result?);
                }
            }
        }

        Ok(())
    }

    pub fn process(&mut self, record: StringRecord) {
        self.inferrence.process(self.selection.select(&record));
        self.buffer.push(record);
    }

    pub fn cast<'a, 'b>(
        &'a self,
        headers: &'b StringRecord,
        record: &'b StringRecord,
    ) -> impl Iterator<Item = Option<(&'a str, Value)>> + 'a
    where
        'b: 'a,
    {
        self.selection
            .select(headers)
            .zip(self.selection.select(record))
            .enumerate()
            .map(|(i, (header, value))| {
                self.inferrence
                    .cast(value, self.inferrence.json_types[i])
                    .map(|json_value| (header, json_value))
            })
    }

    pub fn mutate_attributes(&self, attributes: &mut OmittableAttributes, record: &StringRecord) {
        for ((cell, current_value), json_type) in self
            .selection
            .select(record)
            .zip(attributes.values_mut())
            .zip(self.inferrence.json_types.iter())
        {
            *current_value = self.inferrence.cast(cell, *json_type);
        }
    }

    pub fn records(&self) -> impl Iterator<Item = &StringRecord> {
        self.buffer.iter()
    }

    pub fn types(&self) -> impl Iterator<Item = JSONType> + '_ {
        self.inferrence.json_types.iter().copied()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_inferrence() {
        let mut inferrence = JSONTypeInferrence::new(4, JSONEmptyMode::Null);

        inferrence.process(["1", "1", "", "5"].into_iter());
        inferrence.process(["2", "george", "", "6"].into_iter());
        inferrence.process(["3", "", "", "3.8"].into_iter());

        assert_eq!(
            inferrence.json_types,
            vec![
                JSONType::Integer,
                JSONType::String,
                JSONType::Null,
                JSONType::Float
            ]
        );
    }
}
