use color_eyre::eyre::{Result, eyre};
use serde::{Deserialize, Serialize};
use serde_json::Value;
use std::collections::HashMap;
use std::fs;
use std::path::Path;

/// Jupyter notebook structure
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Notebook {
    pub cells: Vec<Cell>,
    pub metadata: NotebookMetadata,
    pub nbformat: i32,
    pub nbformat_minor: i32,
}

/// Notebook cell
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Cell {
    pub cell_type: String,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub execution_count: Option<i32>,
    pub metadata: HashMap<String, Value>,
    pub source: Vec<String>,
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    pub outputs: Vec<Output>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub id: Option<String>,
}

/// Cell output
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Output {
    pub output_type: String,
    #[serde(flatten)]
    pub data: HashMap<String, Value>,
}

/// Notebook metadata
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct NotebookMetadata {
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub kernelspec: Option<KernelSpec>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub language_info: Option<LanguageInfo>,
    #[serde(flatten)]
    pub extra: HashMap<String, Value>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct KernelSpec {
    pub display_name: String,
    pub language: String,
    pub name: String,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LanguageInfo {
    pub name: String,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub version: Option<String>,
    #[serde(flatten)]
    pub extra: HashMap<String, Value>,
}

/// Read a Jupyter notebook from a file
pub fn read_notebook<P: AsRef<Path>>(path: P) -> Result<Notebook> {
    let content = fs::read_to_string(path)?;
    let notebook: Notebook = serde_json::from_str(&content)?;
    Ok(notebook)
}

/// Write a Jupyter notebook to a file
pub fn write_notebook<P: AsRef<Path>>(path: P, notebook: &Notebook) -> Result<()> {
    let content = serde_json::to_string_pretty(notebook)?;
    fs::write(path, content)?;
    Ok(())
}

/// Edit operations for notebooks
pub enum NotebookEdit {
    /// Replace cell at index with new source
    EditCell { index: usize, source: Vec<String> },
    /// Add a new cell at index
    AddCell { index: usize, cell: Cell },
    /// Delete cell at index
    DeleteCell { index: usize },
}

/// Apply an edit to a notebook
pub fn edit_notebook(notebook: &mut Notebook, edit: NotebookEdit) -> Result<()> {
    match edit {
        NotebookEdit::EditCell { index, source } => {
            if index >= notebook.cells.len() {
                return Err(eyre!("Cell index {} out of bounds (notebook has {} cells)", index, notebook.cells.len()));
            }
            notebook.cells[index].source = source;
        }
        NotebookEdit::AddCell { index, cell } => {
            if index > notebook.cells.len() {
                return Err(eyre!("Cell index {} out of bounds (notebook has {} cells)", index, notebook.cells.len()));
            }
            notebook.cells.insert(index, cell);
        }
        NotebookEdit::DeleteCell { index } => {
            if index >= notebook.cells.len() {
                return Err(eyre!("Cell index {} out of bounds (notebook has {} cells)", index, notebook.cells.len()));
            }
            notebook.cells.remove(index);
        }
    }
    Ok(())
}

/// Create a new code cell
pub fn new_code_cell(source: Vec<String>) -> Cell {
    Cell {
        cell_type: "code".to_string(),
        execution_count: None,
        metadata: HashMap::new(),
        source,
        outputs: Vec::new(),
        id: None,
    }
}

/// Create a new markdown cell
pub fn new_markdown_cell(source: Vec<String>) -> Cell {
    Cell {
        cell_type: "markdown".to_string(),
        execution_count: None,
        metadata: HashMap::new(),
        source,
        outputs: Vec::new(),
        id: None,
    }
}

/// Get cell source as a single string
pub fn cell_source_as_string(cell: &Cell) -> String {
    cell.source.join("")
}

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

    #[test]
    fn cell_source_concatenation() {
        let cell = Cell {
            cell_type: "code".to_string(),
            execution_count: None,
            metadata: HashMap::new(),
            source: vec!["import pandas as pd\n".to_string(), "df = pd.DataFrame()\n".to_string()],
            outputs: Vec::new(),
            id: None,
        };

        let source = cell_source_as_string(&cell);
        assert_eq!(source, "import pandas as pd\ndf = pd.DataFrame()\n");
    }

    #[test]
    fn new_code_cell_creation() {
        let cell = new_code_cell(vec!["print('hello')\n".to_string()]);
        assert_eq!(cell.cell_type, "code");
        assert!(cell.outputs.is_empty());
    }

    #[test]
    fn new_markdown_cell_creation() {
        let cell = new_markdown_cell(vec!["# Title\n".to_string()]);
        assert_eq!(cell.cell_type, "markdown");
    }
}