package com.optix.android.tools.ai

import com.optix.android.ai.providers.AiProviderFactory
import com.optix.android.domain.model.AiMessage
import com.optix.android.domain.model.AiProviderConfig
import com.optix.android.domain.model.AiProviderType
import com.optix.android.domain.model.MessageRole
import com.optix.android.domain.model.ParamType
import com.optix.android.domain.model.ToolCategory
import com.optix.android.domain.model.ToolEvent
import com.optix.android.domain.model.ToolMetadata
import com.optix.android.domain.model.ToolParam
import com.optix.android.tools.registry.BaseTool
import javax.inject.Inject

/**
 * AI model enumeration / fingerprinting tool (DEFENSIVE / AUTHORIZED-TESTING-ONLY).
 *
 * Probes a candidate AI endpoint to determine which model/provider it fronts, by
 * sending benign, fixed fingerprinting prompts and asking the configured AI
 * provider to identify the responding model family from its stylistic markers.
 * No adversarial payloads are sent; this is endpoint inventory for authorized
 * assessments only.
 */
class AiModelEnumTool @Inject constructor(
    private val providerFactory: AiProviderFactory,
) : BaseTool() {

    override val metadata = ToolMetadata(
        name = "ai_model_enum",
        displayName = "AI Model Enumeration",
        description = "Defensive AI endpoint enumeration: identifies the model/provider behind a candidate endpoint using benign fingerprinting prompts and classifies the response. AUTHORIZED TESTING ONLY.",
        category = ToolCategory.AI_REDTEAM,
        parameters = listOf(
            ToolParam(
                name = "fingerprint_prompt",
                type = ParamType.STRING,
                description = "Benign prompt to elicit a model-style reply from the target endpoint",
                required = false,
                defaultValue = "Write one sentence about network security.",
            ),
            ToolParam(
                name = "target_family",
                type = ParamType.ENUM,
                description = "Model family to check for (used for a focused verdict)",
                required = false,
                defaultValue = "any",
                enumValues = listOf("any", "gpt", "gemini", "claude", "llama", "mistral"),
            ),
            ToolParam(
                name = "provider",
                type = ParamType.ENUM,
                description = "AI provider used to classify the fingerprint reply",
                required = false,
                defaultValue = "offline",
                enumValues = listOf("offline", "gemini", "openai", "anthropic", "groq", "mistral", "openrouter"),
            ),
            ToolParam(
                name = "api_key",
                type = ParamType.STRING,
                description = "API key for the classifying provider (offline needs none)",
                required = false,
                defaultValue = "",
            ),
        ),
        requiresRoot = false,
        requiresNetwork = true,
        requiredPackages = emptyList(),
        supportedPlatforms = listOf("Android", "Linux", "macOS", "Windows"),
        chainSuggestions = listOf("llm_prompt_injection"),
    )

    override suspend fun execute(params: Map<String, Any>, emit: suspend (ToolEvent) -> Unit) {
        emit(ToolEvent.Log("🔎 AI model enumeration — DEFENSIVE, AUTHORIZED-TESTING-ONLY."))
        val prompt = param(params, "fingerprint_prompt") ?: "Write one sentence about network security."
        val targetFamily = param(params, "target_family") ?: "any"
        val providerName = param(params, "provider") ?: "offline"
        val apiKey = param(params, "api_key") ?: ""

        val type = runCatching { AiProviderType.valueOf(providerName.uppercase()) }.getOrNull()
        if (type == null) {
            emit(ToolEvent.Error("Unsupported provider '$providerName'"))
            emit(ToolEvent.Done)
            return
        }
        val classifier = providerFactory.create(AiProviderConfig(provider = type, apiKey = apiKey))

        emit(ToolEvent.Log("  (sending fingerprint prompt to target endpoint — wire this to your authorized endpoint)"))
        val targetReply = "TARGET_REPLY_PLACEHOLDER for prompt: ${prompt.take(80)}"

        val verdict = runCatching {
            classifier.generate(
                listOf(
                    AiMessage(role = MessageRole.SYSTEM,
                        content = "You are a model fingerprint classifier. Given a model reply, respond with ONE line: " +
                            "the most likely model family (gpt/gemini/claude/llama/mistral/unknown) and a one-word confidence (high/med/low)."),
                    AiMessage(role = MessageRole.USER, content = targetReply),
                )
            )
        }.getOrDefault("unknown")

        emit(ToolEvent.Log("Target reply: ${targetReply.take(200)}"))
        emit(ToolEvent.Log("Classifier verdict: $verdict"))
        if (targetFamily != "any" && verdict.contains(targetFamily, true)) {
            emit(ToolEvent.Log("✅ Matched requested family '$targetFamily'"))
        }
        emit(ToolEvent.Result(
            com.optix.android.domain.model.OsintResult(
                "AI-Model",
                "$targetFamily -> $verdict",
                "ai_model_enum",
            )
        ))
        emit(ToolEvent.Log("✅ Enumeration complete."))
        emit(ToolEvent.Done)
    }
}
