use crate::core::model::{Category, State, SubTask, Todo};
use crate::core::store::Store;
use anyhow::{Context, Result};
use arrow::array::{Array, BooleanArray, Int64Array, StringArray, UInt64Array};
use arrow::datatypes::{DataType, Field, Schema};
use arrow::record_batch::RecordBatch;
use parquet::arrow::arrow_writer::ArrowWriter;
use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
use std::fs::File;
use std::path::PathBuf;
use std::sync::Arc;

pub struct TruenoStore {
    path: PathBuf,
}

impl TruenoStore {
    pub fn new(path: impl Into<PathBuf>) -> Self {
        Self { path: path.into() }
    }

    fn get_schema() -> Arc<Schema> {
        Arc::new(Schema::new(vec![
            Field::new("id", DataType::UInt64, false),
            Field::new("parent_id", DataType::UInt64, true),
            Field::new("category", DataType::Utf8, false),
            Field::new("title", DataType::Utf8, false),
            Field::new("done", DataType::Boolean, false),
            Field::new("notes", DataType::Utf8, true),
            Field::new("created_at", DataType::Int64, false),
            Field::new("updated_at", DataType::Int64, false),
            Field::new("color", DataType::Utf8, false),
        ]))
    }

    fn flatten_state(state: &State) -> Vec<FlatItem> {
        let mut flat = Vec::new();
        for cat in State::all_categories() {
            let todos = state.get_category(cat);
            for todo in todos {
                flat.push(FlatItem {
                    id: todo.id,
                    parent_id: None,
                    category: cat.as_str().to_string(),
                    title: todo.title.clone(),
                    done: cat == Category::Completed,
                    notes: todo.notes.clone(),
                    created_at: todo.created_at,
                    updated_at: todo.updated_at,
                    color: todo.color.as_str().to_string(),
                });
                Self::flatten_subs(&todo.subtasks, todo.id, cat.as_str(), &mut flat);
            }
        }
        flat
    }

    fn flatten_subs(subs: &[SubTask], parent_id: u64, cat_str: &str, flat: &mut Vec<FlatItem>) {
        for sub in subs {
            flat.push(FlatItem {
                id: sub.id,
                parent_id: Some(parent_id),
                category: cat_str.to_string(),
                title: sub.title.clone(),
                done: sub.done,
                notes: sub.notes.clone(),
                created_at: sub.created_at,
                updated_at: sub.updated_at,
                color: sub.color.as_str().to_string(),
            });
            Self::flatten_subs(&sub.subtasks, sub.id, cat_str, flat);
        }
    }

    fn reconstruct_state(items: Vec<FlatItem>) -> State {
        let mut state = State::new();
        let mut max_id = 1;
        
        // Build map of parent_id -> children
        use std::collections::HashMap;
        let mut children_map: HashMap<Option<u64>, Vec<FlatItem>> = HashMap::new();
        for item in items {
            if item.id >= max_id { max_id = item.id + 1; }
            children_map.entry(item.parent_id).or_default().push(item);
        }

        state.next_id = max_id;

        // Process L1 (parent_id is None)
        if let Some(l1_items) = children_map.remove(&None) {
            for item in l1_items {
                let cat = Category::from_str(&item.category).unwrap_or(Category::Short);
                let todo = Todo {
                    id: item.id,
                    title: item.title,
                    notes: item.notes,
                    created_at: item.created_at,
                    updated_at: item.updated_at,
                    color: crate::core::model::ColorTag::from_str(&item.color),
                    subtasks: Self::reconstruct_subs(item.id, &mut children_map),
                };
                state.get_category_mut(cat).push(todo);
            }
        }

        state
    }

    fn reconstruct_subs(parent_id: u64, map: &mut std::collections::HashMap<Option<u64>, Vec<FlatItem>>) -> Vec<SubTask> {
        let mut subs = Vec::new();
        // Use get if we want to leave the map intact, but remove is fine as we only visit once
        if let Some(children) = map.remove(&Some(parent_id)) {
            for child in children {
                subs.push(SubTask {
                    id: child.id,
                    title: child.title,
                    done: child.done,
                    created_at: child.created_at,
                    updated_at: child.updated_at,
                    color: crate::core::model::ColorTag::from_str(&child.color),
                    notes: child.notes,
                    subtasks: Self::reconstruct_subs(child.id, map),
                });
            }
        }
        subs
    }
}


struct FlatItem {
    id: u64,
    parent_id: Option<u64>,
    category: String,
    title: String,
    done: bool,
    notes: Option<String>,
    created_at: i64,
    updated_at: i64,
    color: String,
}

impl Store for TruenoStore {
    fn load(&self) -> Result<State> {
        if !self.path.exists() {
            return Ok(State::new());
        }
        let file = File::open(&self.path)?;
        let builder = ParquetRecordBatchReaderBuilder::try_new(file)?;
        let mut reader = builder.build()?;
        
        let mut items = Vec::new();
        while let Some(batch_res) = reader.next() {
            let batch = batch_res?;
            let ids = batch.column(0).as_any().downcast_ref::<UInt64Array>().context("id column missing")?;
            let parent_ids = batch.column(1).as_any().downcast_ref::<UInt64Array>().context("parent_id column missing")?;
            let categories = batch.column(2).as_any().downcast_ref::<StringArray>().context("category column missing")?;
            let titles = batch.column(3).as_any().downcast_ref::<StringArray>().context("title column missing")?;
            let dones = batch.column(4).as_any().downcast_ref::<BooleanArray>().context("done column missing")?;
            let notes = batch.column(5).as_any().downcast_ref::<StringArray>().context("notes column missing")?;
            let created_ats = batch.column(6).as_any().downcast_ref::<Int64Array>().context("created_at column missing")?;
            let updated_ats = batch.column(7).as_any().downcast_ref::<Int64Array>().context("updated_at column missing")?;
            let colors = batch.column(8).as_any().downcast_ref::<StringArray>().context("color column missing")?;

            for i in 0..batch.num_rows() {
                items.push(FlatItem {
                    id: ids.value(i),
                    parent_id: if parent_ids.is_null(i) { None } else { Some(parent_ids.value(i)) },
                    category: categories.value(i).to_string(),
                    title: titles.value(i).to_string(),
                    done: dones.value(i),
                    notes: if notes.is_null(i) { None } else { Some(notes.value(i).to_string()) },
                    created_at: created_ats.value(i),
                    updated_at: updated_ats.value(i),
                    color: colors.value(i).to_string(),
                });
            }
        }

        Ok(Self::reconstruct_state(items))
    }

    fn save(&self, state: &State) -> Result<()> {
        let items = Self::flatten_state(state);
        let schema = Self::get_schema();

        let ids: Vec<u64> = items.iter().map(|i| i.id).collect();
        let parent_ids: Vec<Option<u64>> = items.iter().map(|i| i.parent_id).collect();
        let categories: Vec<String> = items.iter().map(|i| i.category.clone()).collect();
        let titles: Vec<String> = items.iter().map(|i| i.title.clone()).collect();
        let dones: Vec<bool> = items.iter().map(|i| i.done).collect();
        let notes: Vec<Option<String>> = items.iter().map(|i| i.notes.clone()).collect();
        let created_ats: Vec<i64> = items.iter().map(|i| i.created_at).collect();
        let updated_ats: Vec<i64> = items.iter().map(|i| i.updated_at).collect();
        let colors: Vec<String> = items.iter().map(|i| i.color.clone()).collect();

        let batch = RecordBatch::try_new(
            schema.clone(),
            vec![
                Arc::new(UInt64Array::from(ids)),
                Arc::new(UInt64Array::from(parent_ids)),
                Arc::new(StringArray::from(categories)),
                Arc::new(StringArray::from(titles)),
                Arc::new(BooleanArray::from(dones)),
                Arc::new(StringArray::from(notes)),
                Arc::new(Int64Array::from(created_ats)),
                Arc::new(Int64Array::from(updated_ats)),
                Arc::new(StringArray::from(colors)),
            ],
        )?;

        let file = File::create(&self.path)
            .with_context(|| format!("Failed to create backup file at {:?}", self.path))?;
        let mut writer = ArrowWriter::try_new(file, schema, None)?;
        writer.write(&batch)?;
        writer.close()?;

        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::core::model::{ColorTag, SubTask, Todo};

    #[test]
    fn test_flatten_state() {
        let mut state = State::new();
        let sub = SubTask {
            id: 2,
            title: "Sub".to_string(),
            done: false,
            created_at: 100,
            updated_at: 100,
            color: ColorTag::Red,
            notes: None,
            subtasks: vec![],
        };
        let todo = Todo {
            id: 1,
            title: "Todo".to_string(),
            notes: Some("Note".to_string()),
            created_at: 50,
            updated_at: 50,
            color: ColorTag::Blue,
            subtasks: vec![sub],
        };
        state.short.push(todo);

        let flat = TruenoStore::flatten_state(&state);
        
        assert_eq!(flat.len(), 2);
        
        let f_todo = flat.iter().find(|i| i.id == 1).unwrap();
        assert_eq!(f_todo.parent_id, None);
        assert_eq!(f_todo.category, "Short");
        
        let f_sub = flat.iter().find(|i| i.id == 2).unwrap();
        assert_eq!(f_sub.parent_id, Some(1));
        assert_eq!(f_sub.category, "Short");
        assert_eq!(f_sub.color, "red");
    }
}

