//
//  ModelCategory.swift
//  oAI
//
//  Category tags for AI models, inferred from model name/id/description.
//
// SPDX-License-Identifier: AGPL-3.0-or-later
// Copyright (C) 2026 Rune Olsen
//
// This file is part of oAI.
//
// oAI is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as
// published by the Free Software Foundation, either version 3 of the
// License, or (at your option) any later version.
//
// oAI is distributed in the hope that it will be useful, but WITHOUT
// ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
// or FITNESS FOR A PARTICULAR PURPOSE. See the GNU Affero General
// Public License for more details.
//
// You should have received a copy of the GNU Affero General Public
// License along with oAI. If not, see <https://www.gnu.org/licenses/>.


import SwiftUI

enum ModelCategory: String, CaseIterable, Codable, Sendable {
    case programming = "Programming"
    case math        = "Math"
    case medical     = "Medical"
    case translation = "Translation"
    case roleplay    = "Roleplay"
    case creative    = "Creative"
    case science     = "Science"
    case finance     = "Finance"
    case legal       = "Legal"

    var color: Color {
        switch self {
        case .programming: return .blue
        case .math:        return .orange
        case .medical:     return .red
        case .translation: return .teal
        case .roleplay:    return .pink
        case .creative:    return .purple
        case .science:     return .green
        case .finance:     return Color(red: 0.75, green: 0.60, blue: 0.0)
        case .legal:       return Color(red: 0.55, green: 0.40, blue: 0.20)
        }
    }

    var systemImage: String {
        switch self {
        case .programming: return "chevron.left.forwardslash.chevron.right"
        case .math:        return "function"
        case .medical:     return "cross.fill"
        case .translation: return "globe"
        case .roleplay:    return "theatermasks.fill"
        case .creative:    return "pencil.and.outline"
        case .science:     return "atom"
        case .finance:     return "chart.line.uptrend.xyaxis"
        case .legal:       return "building.columns.fill"
        }
    }

    // MARK: - Category Inference

    /// Infer categories from a model's name, id, and description.
    static func infer(name: String, id: String, description: String?) -> [ModelCategory] {
        let nameId = (name + " " + id).lowercased()
        let desc   = description?.lowercased() ?? ""
        return allCases.filter { $0.matches(nameId: nameId, desc: desc) }
    }

    private func matches(nameId: String, desc: String) -> Bool {
        if nameKeywords.contains(where: { nameId.contains($0) }) { return true }
        if desc.count > 40 && descKeywords.contains(where: { desc.contains($0) }) { return true }
        return false
    }

    // Patterns matched against lowercased "name + id" string
    private var nameKeywords: [String] {
        switch self {
        case .programming:
            return ["code", "coder", "codex", "codellama", "starcoder", "phind",
                    "codestral", "opencoder", "swe-", "devin-", "wizard-code",
                    "replit-code", "qwen-coder", "deepseek-coder", "devstral",
                    "granite-code", "yi-coder", "artigenz", "wavecoder",
                    "programming", "software-", "cursor-"]
        case .math:
            return ["math", "mathem", "numina", "minerva-math", "wizard-math",
                    "deepseek-math", "qwen-math", "numinamath", "mathstral",
                    "qwq", "internlm-math", "mammoth", "mathcoder", "orion-math",
                    "abel-", "metamath"]
        case .medical:
            return ["medical", "meditron", "med42", "medllama", "biomed",
                    "health-llm", "biosage", "clinicalbert", "pubmedbert",
                    "clinical", "llama-med", "openbiomed", "pmc-llama",
                    "doctorglm", "biolm", "biomistral", "medalpaca",
                    "medpalm", "pharmallm", "mimic"]
        case .translation:
            return ["nllb", "madlad", "-aya-", "seamless", "tower-instruct",
                    "alma-", "bayling", "opus-mt", "m2m-100", "mbart",
                    "translate", "multilingual-", "xglm", "madlad-400"]
        case .roleplay:
            return ["roleplay", "role-play", "mytho", "capybara", "cinematika",
                    "manticore", "weaver-", "noromaid", "airoboros", "toppy",
                    "dolphin", "hermes", "openhermes", "psyfighter",
                    "bluemoon", "midnight", "remm", "rose-20b"]
        case .creative:
            return ["creative-writing", "story-writer", "storyllm", "fimbulvetr",
                    "rp-", "goliath", "lzlv", "mlewd"]
        case .science:
            return ["scibert", "biogpt", "galactica", "science-llm", "scillm",
                    "darwin-", "newton-", "eureka-", "sci-"]
        case .finance:
            return ["fingpt", "finma", "finance-llm", "financellm", "pixiu",
                    "flang-", "alphafin", "bloom-finance", "finbert",
                    "stockgpt", "traderllm"]
        case .legal:
            return ["lawbench", "legalbench", "legalbert", "lawgpt", "legal-llm",
                    "lawyerllm", "legalai", "chatlaw", "jurisllm"]
        }
    }

    // Phrases matched against lowercased description (only when description > 40 chars)
    private var descKeywords: [String] {
        switch self {
        case .programming:
            return ["code generation", "designed for coding", "built for code",
                    "coding-focused", "programming assistant", "specialized for code",
                    "optimized for coding", "coding tasks", "software engineering",
                    "for developers", "code completion", "software development",
                    "coding and", "and coding", "writing code"]
        case .math:
            return ["mathematical reasoning", "math competition", "math olympiad",
                    "designed for math", "theorem proving", "quantitative reasoning",
                    "numerical problem", "math, code", "math and code",
                    "mathematics and", "advanced math", "math tasks",
                    "solving math", "competition math", "math problems"]
        case .medical:
            return ["medical knowledge", "clinical reasoning", "healthcare",
                    "biomedical research", "medical question", "trained on medical",
                    "medical domain", "medical literature", "clinical decision",
                    "health information", "medical text", "medical imaging"]
        case .translation:
            return ["machine translation", "language translation",
                    "multilingual translation", "cross-lingual",
                    "translation tasks", "translation between",
                    "natural language translation"]
        case .roleplay:
            return ["designed for roleplay", "roleplay scenarios",
                    "creative roleplay", "interactive roleplay", "character roleplay",
                    "role-playing", "roleplaying", "uncensored", "nsfw",
                    "adult content", "creative fiction", "interactive story"]
        case .creative:
            return ["creative writing", "storytelling", "narrative generation",
                    "fiction writing", "prose generation", "story writing",
                    "creative text", "story generation", "write stories"]
        case .science:
            return ["scientific literature", "scientific research",
                    "chemistry tasks", "biology research", "physics",
                    "scientific reasoning", "science tasks", "stem tasks",
                    "scientific knowledge"]
        case .finance:
            return ["financial analysis", "quantitative finance",
                    "financial modeling", "market analysis", "financial reasoning",
                    "investment analysis", "economic analysis", "trading"]
        case .legal:
            return ["legal document", "case law", "legal reasoning",
                    "legal text", "law and legal", "legal questions",
                    "legal analysis", "legal research", "contract analysis"]
        }
    }
}