Editor/WeaponImporter/Core/Analysis/WeaponClassifier.cs
#nullable enable annotations

using WeaponImporter.Core.Weapon;

namespace WeaponImporter.Core.Analysis;

/// <summary>
/// Guesses the weapon category from names (file, bones, meshes, clips) and shape (length,
/// bore size, detected parts). The type only picks defaults; users can always override it.
/// </summary>
public static class WeaponClassifier
{
    private static readonly (WeaponType Type, string[] Names)[] NameHints =
    {
        (WeaponType.Revolver, new[] { "revolver", "magnum", "python", "colt", "44", "357", "cylinder", "peacemaker", "rhino" }),
        (WeaponType.Pistol, new[] { "pistol", "handgun", "glock", "deagle", "deserteagle", "eagle", "beretta", "m9", "m1911", "1911", "usp", "p250", "p2000", "fiveseven", "sidearm", "sig", "makarov", "tec9", "cz75" }),
        (WeaponType.Smg, new[] { "smg", "uzi", "mp5", "mp7", "mp9", "p90", "vector", "ump", "mac10", "mac", "thompson", "bizon", "submachine", "pdw", "mp40", "ppsh" }),
        (WeaponType.Shotgun, new[] { "shotgun", "spas", "benelli", "m870", "remington", "nova", "xm1014", "sawedoff", "pump", "doublebarrel", "aa12", "saiga" }),
        (WeaponType.Sniper, new[] { "sniper", "awp", "awm", "barrett", "m82", "l96", "scout", "ssg", "m24", "dragunov", "svd", "kar98", "mosin", "boltaction", "marksman", "dmr", "rifle_sniper" }),
        (WeaponType.Launcher, new[] { "rpg", "launcher", "bazooka", "rocket", "grenadelauncher", "m79", "at4", "javelin", "stinger", "law", "panzerfaust" }),
        (WeaponType.Melee, new[] { "knife", "sword", "katana", "axe", "hatchet", "bat", "crowbar", "machete", "blade", "melee", "club", "wrench", "dagger", "bayonet", "spear", "mace", "pipe" }),
        (WeaponType.Rifle, new[] { "rifle", "ak", "ak47", "ak74", "akm", "m4", "m4a1", "m16", "ar15", "ar", "scar", "famas", "aug", "g36", "galil", "hk416", "assault", "carbine", "fal", "m14" }),
    };

    /// <summary>
    /// Types named in a piece of text. When one alias contains another ("sniperrifle" holds both
    /// "sniper" and "rifle") only the alias the token starts with counts.
    /// </summary>
    private static IEnumerable<WeaponType> NamedTypes(string text)
    {
        var tokens = NameTokens.Split(text);
        var joined = NameTokens.Joined(tokens);
        var hits = new List<(WeaponType Type, string Alias, string Token)>();
        foreach (var (type, names) in NameHints)
            foreach (var n in names)
                foreach (var t in tokens)
                    if (t == n || (n.Length >= 5 && (t.StartsWith(n, StringComparison.Ordinal) || joined.Contains(n, StringComparison.Ordinal))))
                    {
                        hits.Add((type, n, t));
                        break;
                    }
        var kept = hits
            .Where(h => !hits.Any(o => o.Type != h.Type && o.Token == h.Token && o.Token.StartsWith(o.Alias, StringComparison.Ordinal) && !h.Token.StartsWith(h.Alias, StringComparison.Ordinal)))
            .Where(h => !hits.Any(o => o.Type != h.Type && o.Alias.Length > h.Alias.Length && o.Alias.Contains(h.Alias, StringComparison.Ordinal)));
        return kept.Select(h => h.Type).Distinct();
    }

    /// <summary>Type from the file, mesh and bone names alone (used before geometry is trusted).</summary>
    public static WeaponType? FromNames(WeaponAsset asset)
    {
        var scores = new Dictionary<WeaponType, float>();
        var sources = new List<(string, float)> { (asset.Name, 1.6f), (Path.GetFileNameWithoutExtension(asset.SourcePath), 1.6f) };
        sources.AddRange(asset.Mesh.PartNames.Select(n => (n, 0.5f)));
        foreach (var (text, weight) in sources)
            foreach (var type in NamedTypes(text))
                scores[type] = scores.GetValueOrDefault(type) + weight;
        return scores.Count == 0 ? null : scores.OrderByDescending(kv => kv.Value).First().Key;
    }

    public static TypeGuess Classify(WeaponAnalysis a)
    {
        var scores = Enum.GetValues<WeaponType>().ToDictionary(t => t, _ => 0f);
        var reasons = new Dictionary<WeaponType, string>();

        void Score(WeaponType t, float s, string why)
        {
            scores[t] += s;
            if (s > 0f && (!reasons.ContainsKey(t) || s >= 0.5f))
                reasons[t] = why;
        }

        // Names: the file name is the strongest signal, then bones/meshes, then clip names.
        var sources = new List<(string Text, float Weight)> { (a.Asset.Name, 1.6f), (Path.GetFileNameWithoutExtension(a.Asset.SourcePath), 1.6f) };
        sources.AddRange(a.Asset.Mesh.PartNames.Select(n => (n, 0.5f)));
        sources.AddRange(a.Asset.Skeleton.Bones.Where(b => !a.ArmBones.Contains(b.Index)).Select(b => (b.Name, 0.35f)));
        sources.AddRange(a.Asset.Clips.Select(c => (c.Name, 0.25f)));

        // Short aliases ("ar", "ak", "44") must match a whole token.
        foreach (var (text, weight) in sources)
            foreach (var type in NamedTypes(text))
                Score(type, weight, $"name '{text}'");

        // Shape.
        var length = a.Length;
        var bore = a.BoreDiameter;
        var hasMag = a.Part(PartKind.Magazine) is not null;
        var hasTrigger = a.Part(PartKind.Trigger) is not null;
        var hasCylinder = a.Part(PartKind.Cylinder) is not null;
        var hasPump = a.Part(PartKind.Pump) is not null;
        var hasSlide = a.Part(PartKind.Slide) is not null;
        var hasBolt = a.Part(PartKind.Bolt) is not null;
        var hasScope = a.Part(PartKind.Scope) is not null;
        var hasStock = HasStock(a);

        if (hasCylinder) Score(WeaponType.Revolver, 1.2f, "rotating cylinder");
        if (hasSlide) Score(WeaponType.Pistol, 0.8f, "slide");
        if (hasPump) Score(WeaponType.Shotgun, 1.0f, "pump");
        if (hasScope && hasBolt && length > 34f) Score(WeaponType.Sniper, 0.8f, "long, scoped, bolt action");

        if (length is > 4f and < 12f && !hasStock) { Score(WeaponType.Pistol, 0.9f, $"compact ({length:0.#} in)"); Score(WeaponType.Revolver, 0.4f, $"compact ({length:0.#} in)"); }
        if (length is >= 11f and < 27f) Score(WeaponType.Smg, 0.6f, $"{length:0.#} in long");
        if (length is >= 25f and < 42f && hasStock) Score(WeaponType.Rifle, 0.8f, $"{length:0.#} in with a stock");
        if (length is >= 25f and < 42f && !hasStock) Score(WeaponType.Rifle, 0.4f, $"{length:0.#} in long");
        if (length >= 40f) { Score(WeaponType.Sniper, 0.5f, $"very long ({length:0.#} in)"); Score(WeaponType.Launcher, 0.2f, $"very long ({length:0.#} in)"); }
        if (bore > 2.2f && length > 20f) Score(WeaponType.Launcher, 1.0f, $"wide bore ({bore:0.#} in)");
        if (bore is > 0.75f and < 1.6f && length > 24f) Score(WeaponType.Shotgun, 0.3f, "wide barrel");
        if (!hasTrigger && !hasMag && a.Asset.Clips.All(c => !c.Name.Contains("reload", StringComparison.OrdinalIgnoreCase)))
        {
            // Solid, thin shapes without firearm parts read as melee weapons.
            var thin = MathF.Max(a.WeaponBounds.Size.Y, a.WeaponBounds.Size.Z) < length * 0.3f;
            Score(WeaponType.Melee, thin ? 0.7f : 0.35f, "no trigger, magazine or reload");
        }
        if (hasMag && !hasStock && length < 13f) Score(WeaponType.Pistol, 0.3f, "magazine, no stock");

        var ordered = scores.Where(kv => kv.Key != WeaponType.Custom).OrderByDescending(kv => kv.Value).ToArray();
        var best = ordered[0];
        if (best.Value < 0.3f)
            return new TypeGuess(WeaponType.Custom, 0.2f, "not enough evidence", scores);
        var runnerUp = ordered.Length > 1 ? ordered[1].Value : 0f;
        var confidence = Math.Clamp(0.4f + (best.Value - runnerUp) * 0.35f + MathF.Min(best.Value, 2f) * 0.1f, 0.25f, 0.98f);
        return new TypeGuess(best.Key, confidence, reasons.GetValueOrDefault(best.Key, ""), scores);
    }

    /// <summary>A stock: geometry behind the grip region reaching near bore height over several inches.</summary>
    public static bool HasStock(WeaponAnalysis a)
    {
        var p = a.Profile;
        if (a.Length < 16f)
            return false;
        var rearEnd = p.MinX + a.Length * 0.22f;
        var tall = 0;
        var total = 0;
        for (var i = 0; i < p.Count && p.X(i) < rearEnd; i++)
        {
            if (!p.Has(i))
                continue;
            total++;
            if (p.Height(i) > a.BoreDiameter * 2.5f && p.Height(i) > 2.2f)
                tall++;
        }
        return total > 0 && tall > total * 0.5f;
    }
}