AI Agent

interface AIAgent {
  val providerId: String
  val providerName: String
  
  // Initialization
  fun initialize(apiKey: String, context: Context)
  fun reinitializeWithNewModel(apiKey: String, context: Context)
  fun setContext(context: Context)
  fun setProjectData(projectTreeResult: ProjectTreeResult)
  fun clearConversation()
  
  // Code generation
  suspend fun generateCode(
    prompt: String,
    context: String?,
    language: String,
    projectStructure: String?
  ): Result<String>
  
  // Modification tracking
  fun recordModification(
    filePath: String, 
    oldContent: String?, 
    newContent: String, 
    success: Boolean
  )
  fun undoLastModification(): Boolean
  fun getModificationHistory(): List<ModificationAttempt>
  
  // Retry logic
  fun resetAttemptCount()
  fun incrementAttemptCount()
  fun getCurrentAttemptCount(): Int
  fun canRetry(): Boolean
  
  // File operations
  fun writeFile(filePath: String, content: String): FileWriteResult
  fun isInitialized(): Boolean
}

data class ModificationAttempt(
  val timestamp: Long,
  val filePath: String,
  val previousContent: String?,
  val newContent: String,
  val attemptNumber: Int = 0,
  val success: Boolean = false
)

class AIAgentManager(private val context: Context) {
  private val agents = mutableMapOf<String, AIAgent>()
  private var currentAgent: AIAgent? = null
  
  fun registerAgent(agent: AIAgent) {
    agents[agent.providerId] = agent
  }
  
  fun setCurrentAgent(providerId: String, apiKey: String) {
    val agent = agents[providerId] ?: return
    agent.initialize(apiKey, context)
    currentAgent = agent
  }
  
  fun getCurrentAgent(): AIAgent? = currentAgent
  
  suspend fun generateCode(
    prompt: String,
    language: String = "kotlin"
  ): Result<String> {
    val agent = currentAgent ?: return Result.failure(
      IllegalStateException("No agent initialized")
    )
    
    return agent.generateCode(
      prompt = prompt,
      context = getProjectContext(),
      language = language,
      projectStructure = getProjectStructure()
    )
  }
}

object AIAgentRegistry {
  private val agents = mutableMapOf<String, () -> AIAgent>()
  
  fun register(providerId: String, factory: () -> AIAgent) {
    agents[providerId] = factory
  }
  
  fun create(providerId: String): AIAgent? {
    return agents[providerId]?.invoke()
  }
  
  fun getAvailableProviders(): List<String> {
    return agents.keys.toList()
  }
}

// Register agents
fun registerAgents() {
  AIAgentRegistry.register("openai") { OpenAI() }
  AIAgentRegistry.register("anthropic") { Anthropic() }
  AIAgentRegistry.register("google") { Gemini() }
  AIAgentRegistry.register("deepseek") { DeepSeek() }
  AIAgentRegistry.register("grok") { Grok() }
  AIAgentRegistry.register("local") { LocalLLM() }
}

class AICodeCompletionService(
  private val context: Context,
  private val aiAgent: AIAgent
) {
  private val completionCache = mutableMapOf<String, CompletionResult>()
  
  data class CompletionResult(
    val text: String,
    val isMultiLine: Boolean,
    val lineCount: Int
  )
  
  suspend fun getSuggestion(
    currentCode: String,
    cursorPosition: Int,
    language: String = "kotlin"
  ): CompletionResult? = withContext(Dispatchers.IO) {
    if (!aiAgent.isInitialized()) {
      return@withContext null
    }
    
    try {
      val beforeCursor = currentCode.substring(
        0, cursorPosition.coerceAtMost(currentCode.length)
      )
      val afterCursor = currentCode.substring(
        cursorPosition.coerceAtMost(currentCode.length)
      )
      val lastLines = beforeCursor.lines().takeLast(20).joinToString("\n")
      
      // Check cache
      val cacheKey = lastLines.takeLast(150)
      if (completionCache.containsKey(cacheKey)) {
        return@withContext completionCache[cacheKey]
      }
      
      // Build prompt
      val prompt = buildCompletionPrompt(
        lastLines, 
        afterCursor.lines().take(5).joinToString("\n"), 
        language
      )
      
      // Get AI suggestion
      val result = aiAgent.generateCode(
        prompt = prompt,
        context = null,
        language = language,
        projectStructure = null
      )
      
      result.fold(
        onSuccess = { response ->
          val completionResult = extractSuggestion(response)
          if (completionResult != null && completionResult.text.isNotBlank()) {
            completionCache[cacheKey] = completionResult
            return@withContext completionResult
          }
          null
        },
        onFailure = { e ->
          handleException(e)
          null
        }
      )
    } catch (e: Exception) {
      handleException(e)
      null
    }
  }
  
  private fun buildCompletionPrompt(
    before: String,
    after: String,
    language: String
  ): String {
    return """
      Complete the following $language code. Provide only the completion, 
      no explanations.
      
      Code before cursor:
      ```$language
      $before

$after

### Completion UI

```kotlin
class CodeCompletionUI(private val editor: CodeEditor) {
  private var suggestionView: SuggestionView? = null
  
  fun showSuggestion(completion: CompletionResult) {
    suggestionView?.dismiss()
    
    suggestionView = SuggestionView(editor.context).apply {
      setSuggestion(completion.text)
      setOnAcceptListener {
        insertCompletion(completion.text)
      }
      show(editor, editor.cursor.line, editor.cursor.column)
    }
  }
  
  fun hideSuggestion() {
    suggestionView?.dismiss()
    suggestionView = null
  }
  
  private fun insertCompletion(text: String) {
    editor.insertText(text, editor.cursor.line, editor.cursor.column)
    hideSuggestion()
  }
}

data class ProjectTreeResult(
  val structure: String,
  val files: List<FileInfo>,
  val dependencies: List<Dependency>
)

data class FileInfo(
  val path: String,
  val type: FileType,
  val size: Long,
  val content: String?
)

// Set project context
val projectTree = buildProjectTree(workspace)
aiAgent.setProjectData(projectTree)

// Generate code with project awareness
val result = aiAgent.generateCode(
  prompt = "Create a new Repository class for User data",
  context = getCurrentFileContent(),
  language = "kotlin",
  projectStructure = projectTree.structure
)

// Record modification
aiAgent.recordModification(
  filePath = file.path,
  oldContent = originalContent,
  newContent = modifiedContent,
  success = true
)

// View history
val history = aiAgent.getModificationHistory()
history.forEach { attempt ->
  println("File: ${attempt.filePath}")
  println("Attempt: ${attempt.attemptNumber}")
  println("Success: ${attempt.success}")
  println("Time: ${Date(attempt.timestamp)}")
}

// Undo last modification
if (aiAgent.undoLastModification()) {
  println("Modification reverted")
}

class AIExceptions {
  class RateLimitException(message: String) : Exception(message)
  class QuotaExceededException(message: String) : Exception(message)
  class InsufficientBalanceException(message: String) : Exception(message)
  class InvalidApiKeyException(message: String) : Exception(message)
}

private fun handleException(e: Throwable) {
  when (e) {
    is RateLimitException -> {
      showError("Rate limit exceeded. Please wait and try again.")
    }
    is QuotaExceededException -> {
      showError("API quota exceeded. Please check your plan.")
    }
    is InsufficientBalanceException -> {
      showError("Insufficient balance. Please add credits.")
    }
    is InvalidApiKeyException -> {
      showError("Invalid API key. Please check your settings.")
    }
    else -> {
      showError("AI request failed: ${e.message}")
    }
  }
}

suspend fun generateWithRetry(
  prompt: String,
  maxAttempts: Int = 3
): Result<String> {
  aiAgent.resetAttemptCount()
  
  repeat(maxAttempts) { attempt ->
    if (!aiAgent.canRetry()) {
      return Result.failure(
        Exception("Maximum retry attempts reached")
      )
    }
    
    aiAgent.incrementAttemptCount()
    
    val result = aiAgent.generateCode(
      prompt = prompt,
      context = null,
      language = "kotlin",
      projectStructure = null
    )
    
    if (result.isSuccess) {
      return result
    }
    
    // Exponential backoff
    delay(1000L * (attempt + 1))
  }
  
  return Result.failure(Exception("Failed after $maxAttempts attempts"))
}

class ProviderSwitchDialog : DialogFragment() {
  private val providers = listOf(
    "openai" to "OpenAI GPT",
    "anthropic" to "Anthropic Claude",
    "google" to "Google Gemini",
    "deepseek" to "DeepSeek",
    "grok" to "Grok",
    "local" to "Local LLM"
  )
  
  override fun onCreateDialog(savedInstanceState: Bundle?): Dialog {
    return AlertDialog.Builder(requireContext())
      .setTitle("Select AI Provider")
      .setItems(providers.map { it.second }.toTypedArray()) { _, which ->
        switchProvider(providers[which].first)
      }
      .create()
  }
}

class AIPermissionDialog : DialogFragment() {
  override fun onCreateDialog(savedInstanceState: Bundle?): Dialog {
    return AlertDialog.Builder(requireContext())
      .setTitle("AI Features")
      .setMessage(
        "Enable AI-powered code completion and generation? "
        + "This will send code snippets to the selected AI provider."
      )
      .setPositiveButton("Enable") { _, _ ->
        enableAI()
      }
      .setNegativeButton("Cancel", null)
      .create()
  }
}