smart-tree 8.0.1

use super::{ai::AiFormatter, Formatter, PathDisplayMode};
use crate::scanner::{FileNode, TreeStats};
use anyhow::Result;
use serde_json::{json, Value};
use std::io::Write;
use std::path::Path;

pub struct AiJsonFormatter {
    ai_formatter: AiFormatter,
}

impl AiJsonFormatter {
    pub fn new(no_emoji: bool, _path_mode: PathDisplayMode) -> Self {
        Self {
            ai_formatter: AiFormatter::new(no_emoji, _path_mode),
        }
    }
}

impl Formatter for AiJsonFormatter {
    fn format(
        &self,
        writer: &mut dyn Write,
        nodes: &[FileNode],
        stats: &TreeStats,
        root_path: &Path,
    ) -> Result<()> {
        // First get the AI format output as a string
        let mut ai_output = Vec::new();
        self.ai_formatter
            .format(&mut ai_output, nodes, stats, root_path)?;
        let ai_text = String::from_utf8_lossy(&ai_output);

        // Parse the AI output to extract structured data
        let lines = ai_text.lines();
        let mut hex_tree_lines = Vec::new();
        let mut context = None;
        let mut hash = None;
        let mut stats_section = false;
        let mut file_count = 0u64;
        let mut dir_count = 0u64;
        let mut total_size = 0u64;
        let mut file_types = Vec::new();
        let mut large_files = Vec::new();
        let mut date_range = None;

        for line in lines {
            if line == "TREE_HEX_V1:" {
                continue;
            } else if line.starts_with("CONTEXT: ") {
                context = Some(line.strip_prefix("CONTEXT: ").unwrap_or("").to_string());
            } else if line.starts_with("HASH: ") {
                hash = Some(line.strip_prefix("HASH: ").unwrap_or("").to_string());
            } else if line == "STATS:" || line.is_empty() {
                stats_section = true;
            } else if line.starts_with("F:") && stats_section {
                // Parse stats line: F:45 D:12 S:23fc00 (2.3MB)
                let parts: Vec<&str> = line.split_whitespace().collect();
                if let Some(f) = parts.first().and_then(|s| s.strip_prefix("F:")) {
                    file_count = u64::from_str_radix(f, 16).unwrap_or(0);
                }
                if let Some(d) = parts.get(1).and_then(|s| s.strip_prefix("D:")) {
                    dir_count = u64::from_str_radix(d, 16).unwrap_or(0);
                }
                if let Some(s) = parts.get(2).and_then(|s| s.strip_prefix("S:")) {
                    total_size = u64::from_str_radix(s, 16).unwrap_or(0);
                }
            } else if line.starts_with("TYPES: ") && stats_section {
                let types_str = line.strip_prefix("TYPES: ").unwrap_or("");
                for type_entry in types_str.split_whitespace() {
                    if let Some((ext, count_hex)) = type_entry.split_once(':') {
                        if let Ok(count) = u64::from_str_radix(count_hex, 16) {
                            file_types.push(json!({
                                "extension": ext,
                                "count": count
                            }));
                        }
                    }
                }
            } else if line.starts_with("LARGE: ") && stats_section {
                let large_str = line.strip_prefix("LARGE: ").unwrap_or("");
                for file_entry in large_str.split_whitespace() {
                    if let Some((name, size_hex)) = file_entry.split_once(':') {
                        if let Ok(size) = u64::from_str_radix(size_hex, 16) {
                            large_files.push(json!({
                                "name": name,
                                "size": size
                            }));
                        }
                    }
                }
            } else if line.starts_with("DATES: ") && stats_section {
                date_range = Some(line.strip_prefix("DATES: ").unwrap_or("").to_string());
            } else if line == "END_AI" {
                break;
            } else if !stats_section && !line.is_empty() {
                // This is a hex tree line
                hex_tree_lines.push(line.to_string());
            }
        }

        // Build the JSON structure
        let mut json_output = json!({
            "version": "AI_JSON_V1",
            "hash": hash.unwrap_or_else(|| "unknown".to_string()),
            "hex_tree": hex_tree_lines,
            "statistics": {
                "files": file_count,
                "directories": dir_count,
                "total_size": total_size,
                "total_size_mb": format!("{:.1}", total_size as f64 / (1024.0 * 1024.0))
            }
        });

        // Add optional fields
        if let Some(ctx) = context {
            json_output["context"] = Value::String(ctx);
        }

        if !file_types.is_empty() {
            json_output["statistics"]["file_types"] = Value::Array(file_types);
        }

        if !large_files.is_empty() {
            json_output["statistics"]["largest_files"] = Value::Array(large_files);
        }

        if let Some(dates) = date_range {
            json_output["statistics"]["date_range"] = Value::String(dates);
        }

        // Write the JSON output
        writeln!(writer, "{}", serde_json::to_string_pretty(&json_output)?)?;

        Ok(())
    }
}