package com.optix.android.ai

import com.optix.android.ai.providers.AiProviderFactory
import com.optix.android.ai.providers.ProviderEvent
import com.optix.android.ai.providers.ToolSchema
import com.optix.android.ai.providers.ToolSchemaParam
import com.optix.android.domain.model.AiMessage
import com.optix.android.domain.model.AiProviderConfig
import com.optix.android.domain.model.MessageRole
import com.optix.android.domain.model.ToolEvent
import com.optix.android.tools.engine.ToolExecutionEngine
import com.optix.android.tools.registry.BaseTool
import com.optix.android.tools.registry.ToolRegistry
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.flow.Flow
import kotlinx.coroutines.flow.flow
import kotlinx.coroutines.flow.flowOn
import javax.inject.Inject
import javax.inject.Singleton
import timber.log.Timber

// ---------------------------------------------------------------------------
// Events emitted by AiBridge to the ViewModel
// ---------------------------------------------------------------------------

sealed class AiStreamEvent {
    data class Token(val text: String) : AiStreamEvent()
    data class ToolStart(val toolName: String, val args: Map<String, Any>) : AiStreamEvent()
    data class ToolResult(val toolName: String, val result: String) : AiStreamEvent()
    /**
     * Emitted when the assistant produced a valid multi-step chain. The consumer
     * (AI ViewModel) can surface it and forward it to the NETWORK tab.
     */
    data class ToolChain(val chain: com.optix.android.domain.model.ToolChain, val skippedSteps: List<String>) : AiStreamEvent()
    data object Done : AiStreamEvent()
}

// ---------------------------------------------------------------------------
// AiBridge — mirrors Python AIToolBridge.stream_ai_response()
// ---------------------------------------------------------------------------

@Singleton
class AiBridge @Inject constructor(
    private val toolRegistry: ToolRegistry,
    private val toolEngine: ToolExecutionEngine,
    private val providerFactory: AiProviderFactory,
    private val chainParser: ChainParser,
) {
    companion object {
        /** Keep SYSTEM + last 40 messages to avoid context window overflow on long sessions. */
        private const val MAX_HISTORY_MESSAGES = 40
    }
    /**
     * Streams an AI response, handling multi-round tool calling.
     * Emits [AiStreamEvent] tokens, tool start/result events, and Done.
     *
     * History is trimmed to [MAX_HISTORY_MESSAGES] to avoid blowing the context window
     * on long sessions. The SYSTEM message is always preserved at index 0.
     */
    fun streamResponse(
        messages: List<AiMessage>,
        config: AiProviderConfig,
    ): Flow<AiStreamEvent> = flow {
        val provider = providerFactory.create(config)
        val candidateTools = candidateToolsForMessages(messages)
        val schemas = buildToolSchemas(candidateTools)

        // Trim history: keep system message + last MAX_HISTORY_MESSAGES
        val trimmedMessages = trimHistory(messages)
        val currentMessages = trimmedMessages.toMutableList()
        var rounds = 0
        val maxRounds = 4

        while (rounds < maxRounds) {
            var hadToolCall = false

            try {
                provider.stream(currentMessages, schemas).collect { event ->
                    when (event) {
                        is ProviderEvent.Token -> emit(AiStreamEvent.Token(event.text))
                        is ProviderEvent.ToolCall -> {
                            hadToolCall = true
                            emit(AiStreamEvent.ToolStart(event.toolName, event.arguments))
                            currentMessages.add(
                                AiMessage(
                                    role = MessageRole.TOOL_CALL,
                                    content = "",
                                    toolName = event.toolName,
                                    toolCallId = event.id,
                                    toolArguments = event.arguments,
                                )
                            )
                            val result = executeToolCall(event.toolName, event.arguments)
                            // Feed result back to MockBrain for contextual follow-ups
                            if (provider is com.optix.android.ai.providers.OfflineProvider) {
                                provider.mockBrain.recordToolResult(event.toolName, result)
                            }
                            emit(AiStreamEvent.ToolResult(event.toolName, result))
                            currentMessages.add(
                                AiMessage(
                                    role = MessageRole.TOOL_RESULT,
                                    content = result,
                                    toolName = event.toolName,
                                    toolCallId = event.id,
                                )
                            )
                        }
                        is ProviderEvent.Done -> { /* handled by loop condition */ }
                        is ProviderEvent.ToolChain -> {
                            val parsed = chainParser.parse(event.raw)
                            if (parsed.chain != null) {
                                emit(AiStreamEvent.ToolChain(parsed.chain, parsed.skippedSteps))
                            }
                        }
                    }
                }
            } catch (e: Exception) {
                Timber.e(e, "Provider stream failed in round $rounds")
                emit(AiStreamEvent.Token("\n\n⚠️ **Provider Error:** ${e.localizedMessage}"))
                break
            }

            if (!hadToolCall) break
            rounds++
        }

        emit(AiStreamEvent.Done)
    }.flowOn(Dispatchers.IO)

    // ------------------------------------------------------------------
    // Tool schema builder
    // ------------------------------------------------------------------

    fun buildToolSchemas(tools: List<BaseTool>): List<ToolSchema> =
        tools.map { tool ->
            ToolSchema(
                name = tool.metadata.name,
                description = tool.metadata.description,
                parameters = tool.metadata.parameters.associate { p ->
                    p.name to ToolSchemaParam(
                        type = p.type.jsonType(),
                        description = p.description,
                        required = p.required,
                        enum = p.enumValues.takeIf { it.isNotEmpty() },
                    )
                },
            )
        }

    // ------------------------------------------------------------------
    // Tool execution (called during AI tool-calling rounds)
    // ------------------------------------------------------------------

    private suspend fun executeToolCall(toolName: String, args: Map<String, Any>): String {
        val tool = toolRegistry.get(toolName)
            ?: return "Error: Tool '$toolName' not found in registry"

        val results = mutableListOf<String>()
        toolEngine.run(tool, args).collect { event ->
            when (event) {
                is ToolEvent.Result -> results.add(formatResult(event.data))
                is ToolEvent.Error  -> results.add("Error: ${event.message}")
                else -> { /* log events not forwarded to AI */ }
            }
        }
        // Cap at 3000 chars to avoid blowing the context window
        return results.joinToString("\n").take(3000).ifEmpty { "No results returned." }
    }

    private fun formatResult(data: Any): String = when (data) {
        is com.optix.android.domain.model.HostResult -> buildString {
            append("HOST ${data.ip}")
            if (!data.hostname.isNullOrEmpty() && data.hostname != data.ip) append(" (${data.hostname})")
            append(" — ${data.status}")
            if (data.openPorts.isNotEmpty()) append(" | ports: ${data.openPorts.joinToString(", ")}")
            if (data.latencyMs > 0) append(" | ${data.latencyMs.toInt()}ms")
        }
        is com.optix.android.domain.model.PortResult ->
            "${data.port}/${data.protocol}  ${data.state}  ${data.service}"
        is com.optix.android.domain.model.DnsRecord ->
            "${data.type}  ${data.name}  →  ${data.value}${if (data.ttl > 0) "  TTL=${data.ttl}" else ""}"
        is com.optix.android.domain.model.OsintResult ->
            "${data.field}: ${data.value}${if (data.source.isNotEmpty()) " [${data.source}]" else ""}"
        is com.optix.android.domain.model.VulnResult ->
            "${data.cveId} (${data.severity}, CVSS ${data.cvss}): ${data.description}"
        is com.optix.android.domain.model.SubdomainResult ->
            "${data.subdomain}${if (data.ip != null) " → ${data.ip}" else ""}${if (data.status != null) " HTTP ${data.status}" else ""}"
        is com.optix.android.domain.model.DirectoryResult ->
            "${data.path}  HTTP ${data.status}"
        is com.optix.android.domain.model.HashResult ->
            "${data.algorithm}  ${data.hash}${if (data.plaintext != null) "  →  ${data.plaintext}" else ""}"
        is com.optix.android.domain.model.SslResult ->
            "${data.field}: ${data.value}"
        else -> data.toString()
    }

    // ------------------------------------------------------------------
    // History trimming — keeps SYSTEM + last MAX_HISTORY_MESSAGES
    // ------------------------------------------------------------------

    private fun trimHistory(messages: List<AiMessage>): List<AiMessage> {
        val system = messages.filter { it.role == MessageRole.SYSTEM }
        val nonSystem = messages.filter { it.role != MessageRole.SYSTEM }
        val trimmed = if (nonSystem.size > MAX_HISTORY_MESSAGES) {
            nonSystem.takeLast(MAX_HISTORY_MESSAGES)
        } else {
            nonSystem
        }
        return system + trimmed
    }

    // ------------------------------------------------------------------
    // Provide all tools to the LLM context so it can intuitively select them
    // ------------------------------------------------------------------

    private fun candidateToolsForMessages(messages: List<AiMessage>): List<BaseTool> =
        ToolSelectionHeuristics.selectCandidates(toolRegistry.getAll(), messages)
}

private object ToolSelectionHeuristics {
    private const val MAX_CANDIDATE_TOOLS = 12

    private val categoryHints = mapOf(
        com.optix.android.domain.model.ToolCategory.NETWORK_SCANNING to setOf(
            "scan", "scanner", "network", "lan", "wifi", "host", "port", "ports",
            "traceroute", "route", "packet", "sniff", "ssh", "nmap", "masscan", "arp",
            "whois", "osint", "cve", "breach", "email", "phone", "github", "shodan",
            "intel", "metadata", "reputation", "geo", "geolocate", "legion", "topology"
        ),
        com.optix.android.domain.model.ToolCategory.DNS_WEB to setOf(
            "dns", "domain", "subdomain", "http", "https", "url", "web", "header",
            "headers", "cors", "jwt", "sql", "sqli", "xss", "directory", "nikto", "sqlmap"
        ),
        com.optix.android.domain.model.ToolCategory.EXPLOITATION to setOf(
            "exploit", "attack", "brute", "hydra", "metasploit", "msfvenom",
            "deauth", "evil", "spoof", "payload", "shell", "cracker"
        ),
        com.optix.android.domain.model.ToolCategory.PAYLOADS_UTILITIES to setOf(
            "encode", "decode", "hash", "password", "regex", "cipher", "subnet",
            "generator", "utility", "payload", "shell"
        ),
    )

    fun selectCandidates(allTools: List<BaseTool>, messages: List<AiMessage>): List<BaseTool> {
        if (allTools.size <= MAX_CANDIDATE_TOOLS) return allTools

        val lastUserMessage = messages.lastOrNull { it.role == MessageRole.USER }?.content.orEmpty()
        val tokens = tokenize(lastUserMessage)
        if (tokens.isEmpty()) return allTools.take(MAX_CANDIDATE_TOOLS)

        val matchedCategories = categoryHints
            .filterValues { hints -> hints.any(tokens::contains) }
            .keys

        val pool = if (matchedCategories.isEmpty()) {
            allTools
        } else {
            allTools.filter { it.metadata.category in matchedCategories }
        }

        val ranked = pool
            .map { tool -> tool to scoreTool(tool, tokens) }
            .sortedWith(
                compareByDescending<Pair<BaseTool, Int>> { it.second }
                    .thenBy { it.first.metadata.displayName }
            )

        val positiveMatches = ranked.filter { it.second > 0 }.map { it.first }
        if (positiveMatches.isNotEmpty()) return positiveMatches.take(MAX_CANDIDATE_TOOLS)

        if (matchedCategories.isNotEmpty()) return pool.take(MAX_CANDIDATE_TOOLS)

        return ranked.take(MAX_CANDIDATE_TOOLS).map { it.first }
    }

    private fun scoreTool(tool: BaseTool, tokens: Set<String>): Int {
        val metadata = tool.metadata
        var score = 0

        score += weightedOverlap(tokens, tokenize(metadata.name), 5)
        score += weightedOverlap(tokens, tokenize(metadata.displayName), 4)
        score += weightedOverlap(tokens, tokenize(metadata.description), 2)
        score += weightedOverlap(tokens, metadata.parameters.flatMap { tokenize(it.name) }.toSet(), 1)

        if (tokens.any { token -> metadata.name.contains(token) }) score += 2
        return score
    }

    private fun weightedOverlap(tokens: Set<String>, candidates: Set<String>, weight: Int): Int =
        tokens.intersect(candidates).size * weight

    private fun tokenize(input: String): Set<String> =
        input.lowercase()
            .split(Regex("[^a-z0-9]+"))
            .filter { it.length >= 2 }
            .toSet()
}
