HuggingFaceService_Fixed.swift

import Foundation
import UIKit

class HuggingFaceService {
    static let shared = HuggingFaceService()
    
    // Updated API endpoints for Gradio 5.x
    private let baseURL = "https://kylanoconnor-plonk-geolocation.hf.space"
    private let session = URLSession.shared
    
    private init() {}
    
    // MARK: - PLONK Location Analysis
    
    /// Analyze location from image using PLONK model on Hugging Face Space
    func analyzeLocationWithPlonk(
        image: UIImage,
        additionalContext: String? = nil
    ) async throws -> LocationAnalysisData {
        
        // First, discover the working endpoint
        guard let workingEndpoint = await discoverWorkingEndpoint() else {
            print("❌ PLONK: No working endpoint found")
            throw HuggingFaceError.spaceUnavailable
        }
        
        print("✅ PLONK: Using endpoint: \(workingEndpoint)")
        
        // Convert image to base64 with compression
        guard let imageData = image.jpegData(compressionQuality: 0.6) else {
            throw HuggingFaceError.imageProcessingFailed
        }
        
        // If image is too large, compress further
        let finalImageData: Data
        if imageData.count > 1_000_000 { // If larger than 1MB
            guard let smallerImageData = image.jpegData(compressionQuality: 0.3) else {
                throw HuggingFaceError.imageProcessingFailed
            }
            finalImageData = smallerImageData
            print("🗜️ Compressed image to \(finalImageData.count) bytes")
        } else {
            finalImageData = imageData
        }
        
        let base64Image = finalImageData.base64EncodedString()
        
        // Try the request
        let plonkResult = try await makeRequest(to: workingEndpoint, with: base64Image)
        let plonkCoordinates = plonkResult.coordinates
        print("✅ PLONK found coordinates: \(plonkCoordinates.formattedString)")
        
        // Enhance with OpenAI analysis
        return try await enhanceWithOpenAI(
            image: image,
            plonkCoordinates: plonkCoordinates,
            additionalContext: additionalContext
        )
    }
    
    /// Discover the working API endpoint for Gradio 5.x
    private func discoverWorkingEndpoint() async -> String? {
        print("🔍 Discovering working PLONK endpoint...")
        
        // Check if the space is running
        guard let url = URL(string: baseURL) else { return nil }
        
        do {
            let (data, response) = try await session.data(from: url)
            
            guard let httpResponse = response as? HTTPURLResponse,
                  httpResponse.statusCode == 200 else {
                print("❌ Space is not accessible")
                return nil
            }
            
            // Check if it's a Gradio app
            let responseString = String(data: data, encoding: .utf8) ?? ""
            if responseString.contains("gradio") || responseString.contains("Gradio") {
                print("✅ Gradio space is running")
                
                // For Gradio 5.x, try these endpoints in order of preference:
                let possibleEndpoints = [
                    "\(baseURL)/api/predict",      // Standard API endpoint
                    "\(baseURL)/run/predict",      // Alternative endpoint
                    "\(baseURL)/gradio_api/predict" // Legacy endpoint
                ]
                
                // Test each endpoint with a simple POST request
                for endpoint in possibleEndpoints {
                    if let endpointURL = URL(string: endpoint) {
                        var request = URLRequest(url: endpointURL)
                        request.httpMethod = "POST"
                        request.setValue("application/json", forHTTPHeaderField: "Content-Type")
                        
                        // Simple test payload
                        let testPayload = ["data": ["test"]]
                        
                        do {
                            request.httpBody = try JSONSerialization.data(withJSONObject: testPayload)
                            let (_, response) = try await session.data(for: request)
                            if let httpResponse = response as? HTTPURLResponse {
                                print("📡 Testing \(endpoint): \(httpResponse.statusCode)")
                                
                                // 200 = working, 422 = validation error but endpoint exists
                                if httpResponse.statusCode == 200 || httpResponse.statusCode == 422 {
                                    print("✅ Found working endpoint: \(endpoint)")
                                    return endpoint
                                }
                            }
                        } catch {
                            // Continue trying other endpoints
                            continue
                        }
                    }
                }
                
                // If no specific endpoint works, return the default
                return "\(baseURL)/api/predict"
            } else {
                print("⚠️ Space might be loading or not a Gradio app")
                return nil
            }
        } catch {
            print("❌ Error checking space: \(error)")
            return nil
        }
    }
    
    /// Make a request to the PLONK API with proper Gradio 5.x format
    private func makeRequest(to endpoint: String, with base64Image: String) async throws -> LocationAnalysisData {
        guard let url = URL(string: endpoint) else {
            throw HuggingFaceError.invalidResponse
        }
        
        // Gradio 5.x API format for gr.Interface
        let payload: [String: Any] = [
            "data": [base64Image] // Simple array format for gr.Interface
        ]
        
        var request = URLRequest(url: url)
        request.httpMethod = "POST"
        request.setValue("application/json", forHTTPHeaderField: "Content-Type")
        request.timeoutInterval = 60.0 // Longer timeout for model inference
        
        print("🔄 PLONK: Submitting to \(endpoint)")
        print("📦 PLONK: Payload keys: \(payload.keys.joined(separator: ", "))")
        
        do {
            request.httpBody = try JSONSerialization.data(withJSONObject: payload)
        } catch {
            print("❌ PLONK: Failed to serialize request payload")
            throw HuggingFaceError.requestSerializationFailed
        }
        
        let (data, response) = try await session.data(for: request)
        
        guard let httpResponse = response as? HTTPURLResponse else {
            throw HuggingFaceError.invalidResponse
        }
        
        print("📡 PLONK: Response status code: \(httpResponse.statusCode)")
        
        // Log the raw response for debugging
        if let responseString = String(data: data, encoding: .utf8) {
            print("📝 PLONK: Raw response: \(responseString.prefix(500))...")
        }
        
        switch httpResponse.statusCode {
        case 200:
            return try parseResponse(data)
        case 422:
            // Validation error - try alternative format
            print("⚠️ Validation error, trying alternative format...")
            return try await makeRequestAlternativeFormat(to: endpoint, with: base64Image)
        case 404:
            throw HuggingFaceError.serverError(404)
        case 413:
            throw HuggingFaceError.networkError("Image too large")
        case 503:
            throw HuggingFaceError.spaceUnavailable
        default:
            throw HuggingFaceError.serverError(httpResponse.statusCode)
        }
    }
    
    /// Try alternative request format for Gradio 5.x
    private func makeRequestAlternativeFormat(to endpoint: String, with base64Image: String) async throws -> LocationAnalysisData {
        guard let url = URL(string: endpoint) else {
            throw HuggingFaceError.invalidResponse
        }
        
        // Try different formats that Gradio 5.x might expect
        let alternativePayloads: [[String: Any]] = [
            // Format 1: With data URI prefix
            [
                "data": ["data:image/jpeg;base64,\(base64Image)"]
            ],
            // Format 2: With inputs key
            [
                "inputs": [base64Image]
            ],
            // Format 3: Direct base64 with PIL type hint
            [
                "data": [base64Image],
                "type": ["PIL"]
            ],
            // Format 4: Legacy format with data URI
            [
                "data": ["data:image/jpeg;base64,\(base64Image)"],
                "fn_index": 0
            ]
        ]
        
        for (index, payload) in alternativePayloads.enumerated() {
            print("🔄 Trying alternative format \(index + 1)")
            
            var request = URLRequest(url: url)
            request.httpMethod = "POST"
            request.setValue("application/json", forHTTPHeaderField: "Content-Type")
            request.timeoutInterval = 60.0
            
            do {
                request.httpBody = try JSONSerialization.data(withJSONObject: payload)
                let (data, response) = try await session.data(for: request)
                
                guard let httpResponse = response as? HTTPURLResponse else {
                    continue
                }
                
                print("📡 Alternative format \(index + 1) status: \(httpResponse.statusCode)")
                
                // Log response for debugging
                if let responseString = String(data: data, encoding: .utf8) {
                    print("📝 Alternative format \(index + 1) response: \(responseString.prefix(200))...")
                }
                
                if httpResponse.statusCode == 200 {
                    return try parseResponse(data)
                }
                
            } catch {
                print("❌ Alternative format \(index + 1) failed: \(error)")
                continue
            }
        }
        
        throw HuggingFaceError.invalidResponseFormat
    }
    
    // MARK: - Response Parsing
    
    private func parseResponse(_ data: Data) throws -> LocationAnalysisData {
        // Log the response for debugging
        if let responseString = String(data: data, encoding: .utf8) {
            print("📝 PLONK: Full response: \(responseString)")
        }
        
        // Try to parse as Gradio 5.x response format
        do {
            let json = try JSONSerialization.jsonObject(with: data) as? [String: Any]
            
            // Gradio 5.x typically returns {"data": [result], ...}
            if let dataArray = json?["data"] as? [Any],
               let locationString = dataArray.first as? String {
                print("✅ PLONK: Parsed location string: \(locationString)")
                return try parseLocationString(locationString)
            }
            
            // Try other possible formats
            if let result = json?["result"] as? [Any],
               let locationString = result.first as? String {
                return try parseLocationString(locationString)
            }
            
            if let output = json?["output"] as? String {
                return try parseLocationString(output)
            }
            
        } catch {
            print("❌ JSON parsing failed: \(error)")
        }
        
        // Try to parse as direct string
        if let responseString = String(data: data, encoding: .utf8) {
            // Sometimes the response might be a direct string
            if responseString.contains("Latitude:") && responseString.contains("Longitude:") {
                return try parseLocationString(responseString)
            }
        }
        
        throw HuggingFaceError.responseParsingFailed
    }
    
    private func parseLocationString(_ locationString: String) throws -> LocationAnalysisData {
        print("🔍 Parsing location string: \(locationString)")
        
        // Parse the format: "Predicted Location:\nLatitude: 40.748817\nLongitude: -73.985428"
        let lines = locationString.components(separatedBy: .newlines)
        var latitude: Double = 0
        var longitude: Double = 0
        var foundLat = false
        var foundLon = false
        
        for line in lines {
            let trimmedLine = line.trimmingCharacters(in: .whitespaces)
            
            if trimmedLine.contains("Latitude:") {
                let parts = trimmedLine.components(separatedBy: ":")
                if parts.count > 1 {
                    let latString = parts[1].trimmingCharacters(in: .whitespaces)
                    if let parsedLat = Double(latString) {
                        latitude = parsedLat
                        foundLat = true
                        print("✅ Found latitude: \(latitude)")
                    }
                }
            } else if trimmedLine.contains("Longitude:") {
                let parts = trimmedLine.components(separatedBy: ":")
                if parts.count > 1 {
                    let lonString = parts[1].trimmingCharacters(in: .whitespaces)
                    if let parsedLon = Double(lonString) {
                        longitude = parsedLon
                        foundLon = true
                        print("✅ Found longitude: \(longitude)")
                    }
                }
            }
        }
        
        guard foundLat && foundLon else {
            print("❌ Could not parse coordinates from: \(locationString)")
            throw HuggingFaceError.invalidResponseFormat
        }
        
        // Validate coordinates are in reasonable range
        guard latitude >= -90 && latitude <= 90 && longitude >= -180 && longitude <= 180 else {
            print("❌ Invalid coordinates: lat=\(latitude), lon=\(longitude)")
            throw HuggingFaceError.invalidResponseFormat
        }
        
        let locationName = generateLocationName(latitude: latitude, longitude: longitude)
        
        return LocationAnalysisData(
            locationName: locationName,
            coordinates: LocationCoordinates(latitude: latitude, longitude: longitude),
            confidenceRadius: 5.0,
            confidencePercentage: 85.0
        )
    }
    
    /// Enhance PLONK coordinates with OpenAI contextual analysis
    private func enhanceWithOpenAI(
        image: UIImage,
        plonkCoordinates: LocationCoordinates,
        additionalContext: String?
    ) async throws -> LocationAnalysisData {
        
        print("🧠 Enhancing PLONK coordinates with OpenAI analysis...")
        
        // Create enhanced prompt with PLONK coordinates
        let enhancedPrompt = """
        I have used a specialized geolocation AI model (PLONK) to analyze this image and it predicted these coordinates:
        
        **PLONK Prediction:**
        - Latitude: \(plonkCoordinates.latitude)
        - Longitude: \(plonkCoordinates.longitude)
        - Location: \(plonkCoordinates.formattedString)
        
        Now, please analyze this image in detail and provide:
        1. **Location Verification**: Does this coordinate prediction seem accurate?
        2. **Specific Location**: What is the specific landmark or location name?
        3. **Context & Details**: Describe what confirms this location
        4. **Confidence Assessment**: How confident are you? (1-100%)
        
        Additional context: \(additionalContext ?? "None provided")
        """
        
        do {
            let openaiResult = try await OpenAIService.shared.analyzeLocationFromImage(
                image: image,
                additionalContext: enhancedPrompt
            )
            
            return LocationAnalysisData(
                locationName: openaiResult.locationName,
                coordinates: plonkCoordinates, // Use PLONK's precise coordinates
                confidenceRadius: openaiResult.confidenceRadius,
                confidencePercentage: openaiResult.confidencePercentage
            )
            
        } catch {
            print("⚠️ OpenAI enhancement failed: \(error)")
            // Fall back to PLONK-only result
            let fallbackLocationName = generateLocationName(
                latitude: plonkCoordinates.latitude,
                longitude: plonkCoordinates.longitude
            )
            
            return LocationAnalysisData(
                locationName: fallbackLocationName,
                coordinates: plonkCoordinates,
                confidenceRadius: 5.0,
                confidencePercentage: 85.0
            )
        }
    }
    
    private func generateLocationName(latitude: Double, longitude: Double) -> String {
        let latDirection = latitude >= 0 ? "N" : "S"
        let lonDirection = longitude >= 0 ? "E" : "W"
        let region = determineRegion(latitude: latitude, longitude: longitude)
        
        return "\(region) (\(String(format: "%.4f", abs(latitude)))°\(latDirection), \(String(format: "%.4f", abs(longitude)))°\(lonDirection))"
    }
    
    private func determineRegion(latitude: Double, longitude: Double) -> String {
        if latitude >= 24 && latitude <= 50 && longitude >= -125 && longitude <= -65 {
            return "North America"
        } else if latitude >= 35 && latitude <= 70 && longitude >= -10 && longitude <= 40 {
            return "Europe"
        } else if latitude >= -35 && latitude <= 35 && longitude >= -20 && longitude <= 50 {
            return "Africa"
        } else if latitude >= 5 && latitude <= 55 && longitude >= 60 && longitude <= 150 {
            return "Asia"
        } else if latitude >= -45 && latitude <= -10 && longitude >= 110 && longitude <= 180 {
            return "Australia/Oceania"
        } else if latitude >= -60 && latitude <= 15 && longitude >= -85 && longitude <= -30 {
            return "South America"
        } else {
            return "Unknown Region"
        }
    }
    
    // MARK: - Testing & Health Check
    
    /// Test the PLONK API with a simple request
    func testPlonkAPI() async {
        print("🧪 Testing PLONK API...")
        
        guard let workingEndpoint = await discoverWorkingEndpoint() else {
            print("❌ No working endpoint found during test")
            return
        }
        
        print("✅ Found working endpoint for test: \(workingEndpoint)")
        
        // Create a simple test image (red square)
        let size = CGSize(width: 224, height: 224)
        UIGraphicsBeginImageContext(size)
        UIColor.red.setFill()
        UIRectFill(CGRect(origin: .zero, size: size))
        let testImage = UIGraphicsGetImageFromCurrentImageContext()!
        UIGraphicsEndImageContext()
        
        do {
            guard let imageData = testImage.jpegData(compressionQuality: 0.8) else {
                print("❌ Failed to create test image data")
                return
            }
            
            let base64Image = imageData.base64EncodedString()
            let result = try await makeRequest(to: workingEndpoint, with: base64Image)
            print("✅ Test successful! Predicted: \(result.locationName)")
        } catch {
            print("❌ Test failed: \(error)")
        }
    }
    
    /// Check if the Hugging Face Space is available
    func checkSpaceAvailability() async -> Bool {
        guard let url = URL(string: baseURL) else { return false }
        
        do {
            let (_, response) = try await session.data(from: url)
            if let httpResponse = response as? HTTPURLResponse {
                print("🌐 PLONK Space status: \(httpResponse.statusCode)")
                return httpResponse.statusCode == 200
            }
            return false
        } catch {
            print("❌ PLONK Space check failed: \(error)")
            return false
        }
    }
}

// MARK: - Extensions

extension HuggingFaceService {
    
    /// Quick location analysis with error handling
    func quickLocationAnalysis(image: UIImage) async -> LocationAnalysisData? {
        do {
            return try await analyzeLocationWithPlonk(image: image)
        } catch {
            print("🚫 PLONK analysis failed: \(error.localizedDescription)")
            return await plonkInspiredAnalysis(image: image)
        }
    }
    
    /// PLONK-inspired analysis using OpenAI (fallback)
    private func plonkInspiredAnalysis(image: UIImage) async -> LocationAnalysisData? {
        print("🧠 Using PLONK-inspired analysis with OpenAI...")
        
        let plonkStylePrompt = """
        Analyze this image for precise geolocation using visual clues. Act like a state-of-the-art visual geolocation model.
        
        Look for:
        - Architecture styles and building materials
        - Vegetation and landscape features  
        - Street signs, license plates, or text
        - Cultural markers and local customs
        - Geographic features
        - Climate indicators
        - Infrastructure patterns
        
        Provide your best GPS coordinate estimate and explain your reasoning.
        """
        
        do {
            return try await OpenAIService.shared.analyzeLocationFromImage(
                image: image,
                additionalContext: plonkStylePrompt
            )
        } catch {
            print("❌ PLONK-inspired analysis failed: \(error)")
            return nil
        }
    }
}

// MARK: - Error Handling

enum HuggingFaceError: Error, LocalizedError {
    case imageProcessingFailed
    case requestSerializationFailed
    case invalidResponse
    case invalidResponseFormat
    case responseParsingFailed
    case serverError(Int)
    case spaceUnavailable
    case networkError(String)
    case quotaExceeded
    
    var errorDescription: String? {
        switch self {
        case .imageProcessingFailed:
            return "Failed to process the image for analysis"
        case .requestSerializationFailed:
            return "Failed to prepare the request"
        case .invalidResponse:
            return "Invalid response from Hugging Face Space"
        case .invalidResponseFormat:
            return "Unexpected response format from PLONK model"
        case .responseParsingFailed:
            return "Failed to parse the location analysis results"
        case .serverError(let code):
            return "Server error (\(code)). The Hugging Face Space may be unavailable"
        case .spaceUnavailable:
            return "The PLONK model space is currently unavailable"
        case .networkError(let message):
            return "Network error: \(message)"
        case .quotaExceeded:
            return "API quota exceeded. Please try again later"
        }
    }
}