sleepy_agent_ios/SleepyAgent/Services/WebSearchService.swift

import Foundation

// MARK: - Errors

enum WebSearchError: LocalizedError {
    case invalidURL
    case networkError(Error)
    case invalidResponse
    case decodingError(Error)
    case serverError(Int)
    case noResults
    
    var errorDescription: String? {
        switch self {
        case .invalidURL:
            return "Invalid search URL"
        case .networkError(let error):
            return "Network error: \(error.localizedDescription)"
        case .invalidResponse:
            return "Invalid response from server"
        case .decodingError(let error):
            return "Failed to parse response: \(error.localizedDescription)"
        case .serverError(let code):
            return "Server error: HTTP \(code)"
        case .noResults:
            return "No search results found"
        }
    }
}

// MARK: - Data Models

struct SearxngResponse: Codable {
    let query: String
    let results: [SearchResult]
    
    enum CodingKeys: String, CodingKey {
        case query
        case results
    }
    
    init(from decoder: Decoder) throws {
        let container = try decoder.container(keyedBy: CodingKeys.self)
        self.query = try container.decodeIfPresent(String.self, forKey: .query) ?? ""
        self.results = try container.decodeIfPresent([SearchResult].self, forKey: .results) ?? []
    }
}

struct SearchResult: Codable {
    let title: String
    let url: String
    let content: String
    let engine: String?
    
    enum CodingKeys: String, CodingKey {
        case title
        case url
        case content
        case engine
    }
    
    init(from decoder: Decoder) throws {
        let container = try decoder.container(keyedBy: CodingKeys.self)
        self.title = try container.decodeIfPresent(String.self, forKey: .title) ?? ""
        self.url = try container.decodeIfPresent(String.self, forKey: .url) ?? ""
        self.content = try container.decodeIfPresent(String.self, forKey: .content) ?? ""
        self.engine = try container.decodeIfPresent(String.self, forKey: .engine)
    }
}

// MARK: - Configuration

struct WebSearchConfiguration {
    var baseURL: String
    var safeSearch: Int
    var timeout: TimeInterval
    var maxResults: Int
    
    static let `default` = WebSearchConfiguration(
        baseURL: "https://search.sleepy.io",
        safeSearch: 0,
        timeout: 30.0,
        maxResults: 5
    )
}

// MARK: - Web Search Service

/// SearXNG client for web search functionality.
actor WebSearchService {
    
    // MARK: - Properties
    
    private var configuration: WebSearchConfiguration
    private let urlSession: URLSession
    private let decoder: JSONDecoder
    
    /// Current base URL for the SearXNG instance
    var baseURL: String {
        get { configuration.baseURL }
        set { configuration.baseURL = newValue.trimmingCharacters(in: .whitespacesAndNewlines).trimmingCharacters(in: CharacterSet(charactersIn: "/")) }
    }
    
    // MARK: - Initialization
    
    init(
        configuration: WebSearchConfiguration = .default,
        urlSession: URLSession = .shared
    ) {
        self.configuration = configuration
        self.urlSession = urlSession
        self.decoder = JSONDecoder()
        self.decoder.keyDecodingStrategy = .convertFromSnakeCase
    }
    
    /// Updates the base URL for the search endpoint.
    func updateBaseURL(_ url: String) {
        baseURL = url
    }
    
    // MARK: - Search Methods
    
    /// Performs a web search query.
    /// - Parameters:
    ///   - query: The search query string
    ///   - maxResults: Maximum number of results to return (defaults to config)
    /// - Returns: Formatted search results as text for LLM consumption
    /// - Throws: WebSearchError if the search fails
    func search(query: String, maxResults: Int? = nil) async throws -> String {
        let results = try await performSearch(query: query)
        return formatResults(results, query: query, maxResults: maxResults ?? configuration.maxResults)
    }
    
    /// Performs a web search and returns raw results.
    /// - Parameter query: The search query string
    /// - Returns: Array of search results
    /// - Throws: WebSearchError if the search fails
    func searchRaw(query: String) async throws -> [SearchResult] {
        try await performSearch(query: query)
    }
    
    // MARK: - Private Methods
    
    private func performSearch(query: String) async throws -> [SearchResult] {
        guard var urlComponents = URLComponents(string: "\(configuration.baseURL)/search") else {
            throw WebSearchError.invalidURL
        }
        
        urlComponents.queryItems = [
            URLQueryItem(name: "q", value: query),
            URLQueryItem(name: "format", value: "json"),
            URLQueryItem(name: "safesearch", value: String(configuration.safeSearch))
        ]
        
        guard let url = urlComponents.url else {
            throw WebSearchError.invalidURL
        }
        
        var request = URLRequest(url: url)
        request.timeoutInterval = configuration.timeout
        request.setValue("application/json", forHTTPHeaderField: "Accept")
        request.setValue("SleepyAgent-iOS/1.0", forHTTPHeaderField: "User-Agent")
        
        let data: Data
        let response: URLResponse
        
        do {
            (data, response) = try await urlSession.data(for: request)
        } catch {
            throw WebSearchError.networkError(error)
        }
        
        guard let httpResponse = response as? HTTPURLResponse else {
            throw WebSearchError.invalidResponse
        }
        
        guard (200...299).contains(httpResponse.statusCode) else {
            throw WebSearchError.serverError(httpResponse.statusCode)
        }
        
        let searchResponse: SearxngResponse
        do {
            searchResponse = try decoder.decode(SearxngResponse.self, from: data)
        } catch {
            throw WebSearchError.decodingError(error)
        }
        
        return searchResponse.results
    }
    
    private func formatResults(_ results: [SearchResult], query: String, maxResults: Int) -> String {
        guard !results.isEmpty else {
            return "No results found for '\(query)'"
        }
        
        return results.prefix(maxResults).map { result in
            var text = "Title: \(result.title)\n"
            text += "URL: \(result.url)\n"
            text += "Content: \(result.content)"
            return text
        }.joined(separator: "\n\n")
    }
    
    // MARK: - Tool Protocol Support
    
    /// Executes search as a tool call with arguments.
    /// - Parameter arguments: Dictionary containing "query" key
    /// - Returns: Formatted search results
    func execute(arguments: [String: String]) async -> String {
        guard let query = arguments["query"] else {
            return "Error: 'query' parameter is required"
        }
        
        do {
            return try await search(query: query)
        } catch {
            return "Error performing web search: \(error.localizedDescription)"
        }
    }
}

// MARK: - Convenience Extensions

extension WebSearchService {
    /// Quick search with default settings.
    static func quickSearch(query: String, baseURL: String? = nil) async throws -> String {
        var config = WebSearchConfiguration.default
        if let baseURL = baseURL {
            config.baseURL = baseURL
        }
        let service = WebSearchService(configuration: config)
        return try await service.search(query: query)
    }
}