improvise/src/import/wizard.rs

use serde_json::Value;
use anyhow::{anyhow, Result};

use super::analyzer::{FieldKind, FieldProposal, analyze_records, extract_array_at_path, find_array_paths};
use crate::model::Model;
use crate::model::cell::{CellKey, CellValue};

#[derive(Debug, Clone, PartialEq)]
pub enum WizardState {
    Preview,
    SelectArrayPath,
    ReviewProposals,
    NameModel,
    Done,
}

#[derive(Debug)]
pub struct ImportWizard {
    pub state: WizardState,
    pub raw: Value,
    pub array_paths: Vec<String>,
    pub selected_path: String,
    pub records: Vec<Value>,
    pub proposals: Vec<FieldProposal>,
    pub model_name: String,
    pub cursor: usize,
    /// Message to display
    pub message: Option<String>,
}

impl ImportWizard {
    pub fn new(raw: Value) -> Self {
        let array_paths = find_array_paths(&raw);
        let state = if raw.is_array() {
            WizardState::ReviewProposals
        } else if array_paths.len() == 1 {
            WizardState::ReviewProposals
        } else {
            WizardState::SelectArrayPath
        };

        let mut wizard = Self {
            state: WizardState::Preview,
            raw: raw.clone(),
            array_paths: array_paths.clone(),
            selected_path: String::new(),
            records: vec![],
            proposals: vec![],
            model_name: "Imported Model".to_string(),
            cursor: 0,
            message: None,
        };

        // Auto-select if array at root or single path
        if raw.is_array() {
            wizard.select_path("");
        } else if array_paths.len() == 1 {
            let path = array_paths[0].clone();
            wizard.select_path(&path);
        }

        wizard.state = if wizard.records.is_empty() && raw.is_object() {
            WizardState::SelectArrayPath
        } else {
            wizard.advance();
            return wizard;
        };
        wizard
    }

    fn select_path(&mut self, path: &str) {
        self.selected_path = path.to_string();
        if let Some(arr) = extract_array_at_path(&self.raw, path) {
            self.records = arr.clone();
            self.proposals = analyze_records(&self.records);
        }
    }

    pub fn advance(&mut self) {
        self.state = match self.state {
            WizardState::Preview => {
                if self.array_paths.len() <= 1 {
                    WizardState::ReviewProposals
                } else {
                    WizardState::SelectArrayPath
                }
            }
            WizardState::SelectArrayPath => WizardState::ReviewProposals,
            WizardState::ReviewProposals => WizardState::NameModel,
            WizardState::NameModel => WizardState::Done,
            WizardState::Done => WizardState::Done,
        };
        self.cursor = 0;
        self.message = None;
    }

    pub fn confirm_path(&mut self) {
        if self.cursor < self.array_paths.len() {
            let path = self.array_paths[self.cursor].clone();
            self.select_path(&path);
            self.advance();
        }
    }

    pub fn toggle_proposal(&mut self) {
        if self.cursor < self.proposals.len() {
            self.proposals[self.cursor].accepted = !self.proposals[self.cursor].accepted;
        }
    }

    pub fn cycle_proposal_kind(&mut self) {
        if self.cursor < self.proposals.len() {
            let p = &mut self.proposals[self.cursor];
            p.kind = match p.kind {
                FieldKind::Category => FieldKind::Measure,
                FieldKind::Measure => FieldKind::TimeCategory,
                FieldKind::TimeCategory => FieldKind::Label,
                FieldKind::Label => FieldKind::Category,
            };
        }
    }

    pub fn move_cursor(&mut self, delta: i32) {
        let len = match self.state {
            WizardState::SelectArrayPath => self.array_paths.len(),
            WizardState::ReviewProposals => self.proposals.len(),
            _ => 0,
        };
        if len == 0 { return; }
        if delta > 0 {
            self.cursor = (self.cursor + 1).min(len - 1);
        } else if self.cursor > 0 {
            self.cursor -= 1;
        }
    }

    pub fn push_name_char(&mut self, c: char) {
        self.model_name.push(c);
    }

    pub fn pop_name_char(&mut self) {
        self.model_name.pop();
    }

    pub fn build_model(&self) -> Result<Model> {
        let mut model = Model::new(self.model_name.clone());

        // Collect categories and measures from accepted proposals
        let categories: Vec<&FieldProposal> = self.proposals.iter()
            .filter(|p| p.accepted && matches!(p.kind, FieldKind::Category | FieldKind::TimeCategory))
            .collect();
        let measures: Vec<&FieldProposal> = self.proposals.iter()
            .filter(|p| p.accepted && p.kind == FieldKind::Measure)
            .collect();

        if categories.is_empty() {
            return Err(anyhow!("At least one category must be accepted"));
        }

        // Add categories
        for cat_proposal in &categories {
            model.add_category(&cat_proposal.field)?;
            if let Some(cat) = model.category_mut(&cat_proposal.field) {
                for val in &cat_proposal.distinct_values {
                    cat.add_item(val);
                }
            }
        }

        // If there are measures, add a "Measure" category
        if !measures.is_empty() {
            model.add_category("Measure")?;
            if let Some(cat) = model.category_mut("Measure") {
                for m in &measures {
                    cat.add_item(&m.field);
                }
            }
        }

        // Import records as cells
        for record in &self.records {
            if let Value::Object(map) = record {
                // Build base coordinate from category fields
                let mut coords: Vec<(String, String)> = vec![];
                let mut valid = true;

                for cat_proposal in &categories {
                    let val = map.get(&cat_proposal.field)
                        .and_then(|v| v.as_str())
                        .map(|s| s.to_string())
                        .or_else(|| map.get(&cat_proposal.field).map(|v| v.to_string()));

                    if let Some(v) = val {
                        // Ensure item exists
                        if let Some(cat) = model.category_mut(&cat_proposal.field) {
                            cat.add_item(&v);
                        }
                        coords.push((cat_proposal.field.clone(), v));
                    } else {
                        valid = false;
                        break;
                    }
                }

                if !valid { continue; }

                // Add each measure as a cell
                for measure in &measures {
                    if let Some(val) = map.get(&measure.field).and_then(|v| v.as_f64()) {
                        let mut cell_coords = coords.clone();
                        if !measures.is_empty() {
                            cell_coords.push(("Measure".to_string(), measure.field.clone()));
                        }
                        let key = CellKey::new(cell_coords);
                        model.set_cell(key, CellValue::Number(val));
                    }
                }
            }
        }

        Ok(model)
    }

    pub fn preview_summary(&self) -> String {
        match &self.raw {
            Value::Array(arr) => {
                format!(
                    "Array of {} records. Sample keys: {}",
                    arr.len(),
                    arr.first()
                        .and_then(|r| r.as_object())
                        .map(|m| m.keys().take(5).cloned().collect::<Vec<_>>().join(", "))
                        .unwrap_or_default()
                )
            }
            Value::Object(map) => {
                format!(
                    "Object with {} top-level keys: {}",
                    map.len(),
                    map.keys().take(10).cloned().collect::<Vec<_>>().join(", ")
                )
            }
            _ => "Unknown JSON structure".to_string(),
        }
    }
}