sherlock-io 1.2.0

use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::fs;
use anyhow::{Context, Result};

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LanguageInfo {
    pub name: String,
    pub extensions: Vec<String>,
    pub patterns: Vec<String>,
    pub category: LanguageCategory,
    pub color: String,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum LanguageCategory {
    Programming,
    Markup,
    Stylesheet,
    Data,
    Configuration,
    Documentation,
    Other,
}

#[derive(Debug, Deserialize)]
struct LanguageConfig {
    languages: HashMap<String, LanguageConfigEntry>,
    special_files: HashMap<String, String>,
}

#[derive(Debug, Deserialize)]
struct LanguageConfigEntry {
    extensions: Vec<String>,
    patterns: Vec<String>,
    category: LanguageCategory,
    color: String,
}

pub struct LanguageRegistry {
    languages: HashMap<String, LanguageInfo>,
    extension_map: HashMap<String, Vec<String>>, // Multiple languages per extension
    special_files: HashMap<String, String>,
}

impl LanguageRegistry {
    pub fn new() -> Result<Self> {
        let mut registry = Self {
            languages: HashMap::new(),
            extension_map: HashMap::new(),
            special_files: HashMap::new(),
        };
        registry.load_languages_from_file()?;
        Ok(registry)
    }

    pub fn from_default() -> Self {
        let mut registry = Self {
            languages: HashMap::new(),
            extension_map: HashMap::new(),
            special_files: HashMap::new(),
        };
        registry.initialize_default_languages();
        registry
    }

    pub fn get_language_by_extension(&self, extension: &str) -> Option<&LanguageInfo> {
        // If only one language uses this extension, return it directly
        if let Some(lang_names) = self.extension_map.get(extension) {
            if lang_names.len() == 1 {
                return self.languages.get(&lang_names[0]);
            }
            // Multiple languages use this extension - return the first one as fallback
            // The caller should use content-based detection for disambiguation
            return self.languages.get(&lang_names[0]);
        }
        None
    }

    /// Get all languages that use a specific extension
    pub fn get_languages_by_extension(&self, extension: &str) -> Vec<&LanguageInfo> {
        if let Some(lang_names) = self.extension_map.get(extension) {
            lang_names.iter()
                .filter_map(|name| self.languages.get(name))
                .collect()
        } else {
            Vec::new()
        }
    }

    pub fn get_language_by_name(&self, name: &str) -> Option<&LanguageInfo> {
        self.languages.get(name)
    }

    pub fn get_all_languages(&self) -> &HashMap<String, LanguageInfo> {
        &self.languages
    }

    pub fn detect_by_content(&self, content: &str, filename: &str) -> Option<&LanguageInfo> {
        // Check special files first
        if let Some(language_name) = self.special_files.get(filename) {
            return self.get_language_by_name(language_name);
        }

        // Shebang detection (highest priority)
        if let Some(first_line) = content.lines().next() {
            if first_line.starts_with("#!") {
                if first_line.contains("python") {
                    return self.get_language_by_name("Python");
                }
                if first_line.contains("node") || first_line.contains("javascript") {
                    return self.get_language_by_name("JavaScript");
                }
                if first_line.contains("bash") || first_line.contains("sh") {
                    return self.get_language_by_name("Shell");
                }
            }
        }

        // Advanced pattern-based detection with improved scoring
        let mut language_scores: Vec<(&LanguageInfo, f64)> = Vec::new();
        
        for (_, lang) in &self.languages {
            let mut score = 0.0;
            let mut pattern_matches = 0;
            let mut unique_patterns_matched = std::collections::HashSet::new();
            
            for pattern in &lang.patterns {
                if content.contains(pattern) {
                    pattern_matches += 1;
                    unique_patterns_matched.insert(pattern);
                    
                    // Calculate pattern specificity score
                    let mut pattern_score = self.calculate_pattern_specificity(pattern);
                    
                    // Bonus for patterns at the beginning of content (more significant)
                    if content.starts_with(pattern) {
                        pattern_score *= 2.0;
                    }
                    
                    // Bonus for patterns that appear multiple times
                    let occurrences = content.matches(pattern).count();
                    if occurrences > 1 {
                        pattern_score *= 1.0 + (occurrences as f64 * 0.1);
                    }
                    
                    score += pattern_score;
                }
            }
            
            if pattern_matches > 0 {
                // Bonus for matching multiple different patterns (stronger language signal)
                let diversity_bonus = if unique_patterns_matched.len() > 1 {
                    1.0 + (unique_patterns_matched.len() as f64 * 0.3)
                } else {
                    1.0
                };
                
                // Normalize by pattern count to favor languages with higher match ratio
                let match_ratio = pattern_matches as f64 / lang.patterns.len() as f64;
                let final_score = score * diversity_bonus * match_ratio;
                
                language_scores.push((lang, final_score));
            }
        }
        
        // Return the language with the highest score
        language_scores.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
        language_scores.first().map(|(lang, _)| *lang)
    }

    /// Detect language by content among a specific set of candidate languages
    pub fn detect_by_content_among_candidates<'a>(&'a self, content: &str, filename: &str, candidates: &[&'a LanguageInfo]) -> Option<&'a LanguageInfo> {
        // Check special files first
        if let Some(language_name) = self.special_files.get(filename) {
            // Check if the special file language is among candidates
            if let Some(lang_info) = self.get_language_by_name(language_name) {
                if candidates.iter().any(|candidate| candidate.name == lang_info.name) {
                    return Some(lang_info);
                }
            }
        }

        // Shebang detection (highest priority)
        if let Some(first_line) = content.lines().next() {
            if first_line.starts_with("#!") {
                let shebang_languages = [
                    ("python", "Python"),
                    ("node", "JavaScript"),
                    ("javascript", "JavaScript"),
                    ("bash", "Shell"),
                    ("sh", "Shell"),
                ];
                
                for (shebang_pattern, lang_name) in &shebang_languages {
                    if first_line.contains(shebang_pattern) {
                        if let Some(lang_info) = self.get_language_by_name(lang_name) {
                            if candidates.iter().any(|candidate| candidate.name == lang_info.name) {
                                return Some(lang_info);
                            }
                        }
                    }
                }
            }
        }

        // Advanced pattern-based detection with improved scoring (only among candidates)
        let mut language_scores: Vec<(&LanguageInfo, f64)> = Vec::new();
        
        for &candidate in candidates {
            let mut score = 0.0;
            let mut pattern_matches = 0;
            let mut unique_patterns_matched = std::collections::HashSet::new();
            
            for pattern in &candidate.patterns {
                if content.contains(pattern) {
                    pattern_matches += 1;
                    unique_patterns_matched.insert(pattern);
                    
                    // Calculate pattern specificity score
                    let mut pattern_score = self.calculate_pattern_specificity(pattern);
                    
                    // Bonus for patterns at the beginning of content (more significant)
                    if content.starts_with(pattern) {
                        pattern_score *= 2.0;
                    }
                    
                    // Bonus for patterns that appear multiple times
                    let occurrences = content.matches(pattern).count();
                    if occurrences > 1 {
                        pattern_score *= 1.0 + (occurrences as f64 * 0.1);
                    }
                    
                    score += pattern_score;
                }
            }
            
            if pattern_matches > 0 {
                // Bonus for matching multiple different patterns (stronger language signal)
                let diversity_bonus = if unique_patterns_matched.len() > 1 {
                    1.0 + (unique_patterns_matched.len() as f64 * 0.3)
                } else {
                    1.0
                };
                
                // Normalize by pattern count to favor languages with higher match ratio
                let match_ratio = pattern_matches as f64 / candidate.patterns.len() as f64;
                let final_score = score * diversity_bonus * match_ratio;
                
                language_scores.push((candidate, final_score));
            }
        }
        
        // Return the language with the highest score
        language_scores.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
        language_scores.first().map(|(lang, _)| *lang)
    }

    /// Calculate pattern specificity score based on uniqueness and characteristics
    fn calculate_pattern_specificity(&self, pattern: &str) -> f64 {
        // Count how many languages use this exact pattern
        let mut usage_count = 0;
        for (_, lang) in &self.languages {
            if lang.patterns.contains(&pattern.to_string()) {
                usage_count += 1;
            }
        }
        
        // Base score inversely proportional to usage (rarer patterns = higher score)
        let rarity_score = if usage_count > 0 {
            10.0 / usage_count as f64
        } else {
            1.0
        };
        
        // Length bonus (longer patterns are generally more specific)
        let length_score = pattern.len() as f64;
        
        // Special character bonus (patterns with special chars are more specific)
        let special_char_bonus = pattern.chars()
            .filter(|c| !c.is_alphanumeric() && !c.is_whitespace())
            .count() as f64 * 0.5;
        
        // Keyword specificity bonus
        let keyword_bonus = if self.is_language_specific_keyword(pattern) {
            2.0
        } else {
            1.0
        };
        
        rarity_score + length_score + special_char_bonus + keyword_bonus
    }
    
    /// Check if a pattern is a language-specific keyword
    fn is_language_specific_keyword(&self, pattern: &str) -> bool {
        // Highly specific patterns that are unique to certain languages
        let specific_keywords = [
            "package ", "<?php", "#!/", "fn main(", "public class", "def __init__",
            "import React", "export default", "namespace", "using System",
            "#include <", "std::", "func main()", "class Program", "module Main",
            "defmodule", "(defn", "pragma solidity", "FROM ", "RUN ", "COPY ",
        ];
        
        specific_keywords.iter().any(|&keyword| pattern.contains(keyword))
    }

    fn load_languages_from_file(&mut self) -> Result<()> {
        let yaml_content = fs::read_to_string("languages.yaml")
            .context("Failed to read languages.yaml file")?;
        
        let config: LanguageConfig = serde_yaml::from_str(&yaml_content)
            .context("Failed to parse languages.yaml")?;

        // Load special files mapping
        self.special_files = config.special_files;

        // Convert and load languages
        for (name, entry) in config.languages {
            let lang_info = LanguageInfo {
                name: name.clone(),
                extensions: entry.extensions.clone(),
                patterns: entry.patterns,
                category: entry.category,
                color: entry.color,
            };

            // Map extensions to language names (multiple languages can share extensions)
            for ext in &entry.extensions {
                self.extension_map
                    .entry(ext.clone())
                    .or_insert_with(Vec::new)
                    .push(name.clone());
            }

            self.languages.insert(name, lang_info);
        }

        Ok(())
    }

    fn initialize_default_languages(&mut self) {
        let languages = vec![
            // Programming Languages
            ("Rust", vec!["rs"], vec!["fn main(", "use std::", "impl"], LanguageCategory::Programming, "#CE422B"),
            ("Python", vec!["py", "pyw"], vec!["def ", "import ", "from "], LanguageCategory::Programming, "#3776AB"),
            ("JavaScript", vec!["js", "mjs"], vec!["function", "const ", "let "], LanguageCategory::Programming, "#F7DF1E"),
            ("TypeScript", vec!["ts"], vec!["interface ", "type ", "export "], LanguageCategory::Programming, "#3178C6"),
            ("Go", vec!["go"], vec!["package ", "func ", "import "], LanguageCategory::Programming, "#00ADD8"),
            ("Java", vec!["java"], vec!["public class", "import java", "public static"], LanguageCategory::Programming, "#ED8B00"),
            ("C", vec!["c", "h"], vec!["#include", "int main", "printf"], LanguageCategory::Programming, "#A8B9CC"),
            ("C++", vec!["cpp", "cc", "cxx", "hpp"], vec!["#include", "std::", "namespace"], LanguageCategory::Programming, "#00599C"),
            ("C#", vec!["cs"], vec!["using System", "namespace", "public class"], LanguageCategory::Programming, "#239120"),
            ("PHP", vec!["php"], vec!["<?php", "function", "$"], LanguageCategory::Programming, "#777BB4"),
            ("Ruby", vec!["rb"], vec!["def ", "class ", "require"], LanguageCategory::Programming, "#CC342D"),
            ("Swift", vec!["swift"], vec!["import ", "func ", "var "], LanguageCategory::Programming, "#FA7343"),
            ("Kotlin", vec!["kt"], vec!["fun ", "class ", "import "], LanguageCategory::Programming, "#7F52FF"),
            ("Scala", vec!["scala"], vec!["object ", "class ", "def "], LanguageCategory::Programming, "#DC322F"),
            ("Shell", vec!["sh", "bash", "zsh"], vec!["#!/bin/", "echo ", "if ["], LanguageCategory::Programming, "#89E051"),
            
            // Web Technologies
            ("HTML", vec!["html", "htm"], vec!["<!DOCTYPE", "<html", "<body"], LanguageCategory::Markup, "#E34F26"),
            ("CSS", vec!["css"], vec!["{", "}", "color:"], LanguageCategory::Stylesheet, "#1572B6"),
            ("SCSS", vec!["scss"], vec!["$", "@import", "@mixin"], LanguageCategory::Stylesheet, "#CF649A"),
            ("Vue", vec!["vue"], vec!["<template>", "<script>", "export default"], LanguageCategory::Programming, "#4FC08D"),
            ("React JSX", vec!["jsx"], vec!["import React", "export default", "className="], LanguageCategory::Programming, "#61DAFB"),
            ("React TSX", vec!["tsx"], vec!["import React", "export default", "interface"], LanguageCategory::Programming, "#61DAFB"),
            
            // Data & Config
            ("JSON", vec!["json"], vec!["{", "}", "\""], LanguageCategory::Data, "#000000"),
            ("YAML", vec!["yml", "yaml"], vec!["---", ":", "- "], LanguageCategory::Data, "#CB171E"),
            ("TOML", vec!["toml"], vec!["[", "]", "="], LanguageCategory::Configuration, "#9C4221"),
            ("XML", vec!["xml"], vec!["<?xml", "<", "/>"], LanguageCategory::Data, "#0060AC"),
            ("SQL", vec!["sql"], vec!["SELECT", "FROM", "WHERE"], LanguageCategory::Data, "#4479A1"),
            
            // Documentation
            ("Markdown", vec!["md"], vec!["# ", "## ", "```"], LanguageCategory::Documentation, "#083FA1"),
            ("Text", vec!["txt"], vec![], LanguageCategory::Documentation, "#000000"),
            
            // Other
            ("Dockerfile", vec!["dockerfile"], vec!["FROM", "RUN", "COPY"], LanguageCategory::Configuration, "#384D54"),
            ("Makefile", vec!["makefile", "mk"], vec![":", "\t", ".PHONY"], LanguageCategory::Configuration, "#427819"),
        ];

        for (name, extensions, patterns, category, color) in languages {
            let lang_info = LanguageInfo {
                name: name.to_string(),
                extensions: extensions.iter().map(|s| s.to_string()).collect(),
                patterns: patterns.iter().map(|s| s.to_string()).collect(),
                category,
                color: color.to_string(),
            };

            // Map extensions to language names (multiple languages can share extensions)
            for ext in &lang_info.extensions {
                self.extension_map
                    .entry(ext.clone())
                    .or_insert_with(Vec::new)
                    .push(name.to_string());
            }

            self.languages.insert(name.to_string(), lang_info);
        }
    }
}