EnhancedTerminal/AIManager.swift

import Foundation

// MARK: - AI Service Protocol
protocol AIService {
    func generateCommand(from prompt: String, completion: @escaping (Result<String, Error>) -> Void)
}

// MARK: - OpenAI Service
class OpenAIService: AIService {
    private let apiKey: String
    
    init(apiKey: String) {
        self.apiKey = apiKey
    }
    
    func generateCommand(from prompt: String, completion: @escaping (Result<String, Error>) -> Void) {
        guard let url = URL(string: "https://api.openai.com/v1/chat/completions") else {
            completion(.failure(NSError(domain: "OpenAI", code: -1, userInfo: [NSLocalizedDescriptionKey: "Invalid URL"])))
            return
        }
        
        var request = URLRequest(url: url)
        request.httpMethod = "POST"
        request.setValue("Bearer \(apiKey)", forHTTPHeaderField: "Authorization")
        request.setValue("application/json", forHTTPHeaderField: "Content-Type")
        
        let body: [String: Any] = [
            "model": "gpt-3.5-turbo",
            "messages": [
                ["role": "system", "content": "You are a terminal command expert. Generate only the command, no explanations. Respond with just the command."],
                ["role": "user", "content": prompt]
            ],
            "max_tokens": 150,
            "temperature": 0.7
        ]
        
        request.httpBody = try? JSONSerialization.data(withJSONObject: body)
        
        URLSession.shared.dataTask(with: request) { data, response, error in
            if let error = error {
                completion(.failure(error))
                return
            }
            
            guard let data = data,
                  let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any],
                  let choices = json["choices"] as? [[String: Any]],
                  let firstChoice = choices.first,
                  let message = firstChoice["message"] as? [String: Any],
                  let content = message["content"] as? String else {
                completion(.failure(NSError(domain: "OpenAI", code: -1, userInfo: [NSLocalizedDescriptionKey: "Invalid response"])))
                return
            }
            
            completion(.success(content.trimmingCharacters(in: .whitespacesAndNewlines)))
        }.resume()
    }
}

// MARK: - Ollama Service
class OllamaService: AIService {
    private let baseURL: String
    private let model: String
    
    init(baseURL: String, model: String) {
        self.baseURL = baseURL
        self.model = model
    }
    
    func generateCommand(from prompt: String, completion: @escaping (Result<String, Error>) -> Void) {
        guard let url = URL(string: "\(baseURL)/api/generate") else {
            completion(.failure(NSError(domain: "Ollama", code: -1, userInfo: [NSLocalizedDescriptionKey: "Invalid URL"])))
            return
        }
        
        var request = URLRequest(url: url)
        request.httpMethod = "POST"
        request.setValue("application/json", forHTTPHeaderField: "Content-Type")
        request.timeoutInterval = 30
        
        let fullPrompt = "Generate only the terminal command for: \(prompt). Respond with just the command, no explanation."
        
        let body: [String: Any] = [
            "model": model,
            "prompt": fullPrompt,
            "stream": false
        ]
        
        request.httpBody = try? JSONSerialization.data(withJSONObject: body)
        
        URLSession.shared.dataTask(with: request) { data, response, error in
            if let error = error {
                completion(.failure(error))
                return
            }
            
            guard let data = data,
                  let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any],
                  let response = json["response"] as? String else {
                completion(.failure(NSError(domain: "Ollama", code: -1, userInfo: [NSLocalizedDescriptionKey: "Invalid response"])))
                return
            }
            
            completion(.success(response.trimmingCharacters(in: .whitespacesAndNewlines)))
        }.resume()
    }
}

// MARK: - AI Manager
class AIManager {
    static let shared = AIManager()
    private var service: AIService?
    
    init() {
        updateService()
    }
    
    func updateService() {
        let settings = TerminalSettings.shared
        switch settings.aiProvider {
        case .openai:
            if !settings.openAIKey.isEmpty {
                service = OpenAIService(apiKey: settings.openAIKey)
            } else {
                service = nil
            }
        case .ollama:
            service = OllamaService(baseURL: settings.ollamaURL, model: settings.ollamaModel)
        case .none:
            service = nil
        }
    }
    
    func isAvailable() -> Bool {
        return service != nil
    }
    
    func generateCommand(from prompt: String, completion: @escaping (Result<String, Error>) -> Void) {
        guard let service = service else {
            completion(.failure(NSError(domain: "AI", code: -1, userInfo: [NSLocalizedDescriptionKey: "No AI service configured"])))
            return
        }
        service.generateCommand(from: prompt, completion: completion)
    }
}