import Foundation

/// Extracts field values from partial/incomplete JSON strings during streaming.
/// Handles cases where JSON is incomplete or malformed.
public struct PartialJsonExtractor {
    public init() {}

    /// Extracts the value of a specific field from partial JSON.
    ///
    /// Example: `{"message":"Hello"}` -> "Hello"
    /// Example: `{"message":"Hello" (incomplete)` -> "Hello"
    /// Example: `{"message":"He` -> "He"
    ///
    /// - Parameters:
    ///   - partialJson: The incomplete JSON string
    ///   - fieldName: The field name to extract (e.g., "message")
    /// - Returns: The extracted field value, or nil if not found or not extractable
    public func extractField(
        partialJson: String,
        fieldName: String
    ) -> String? {
        // Look for the field name pattern: "fieldName":"
        let fieldPattern = "\"\(fieldName)\":"
        guard let fieldRange = partialJson.range(of: fieldPattern) else {
            return nil
        }

        // Find the start of the value after the field name
        let valueStartIndex = partialJson.index(fieldRange.upperBound, offsetBy: 0)
        guard valueStartIndex < partialJson.endIndex else {
            return nil
        }

        // Skip whitespace after the colon
        var currentIndex = valueStartIndex
        while currentIndex < partialJson.endIndex && partialJson[currentIndex].isWhitespace {
            currentIndex = partialJson.index(after: currentIndex)
        }
        guard currentIndex < partialJson.endIndex else {
            return nil
        }

        if partialJson[currentIndex] == "\"" {
            return extractStringValue(json: partialJson, startIndex: partialJson.index(after: currentIndex))
        } else {
            return extractNonStringValue(json: partialJson, startIndex: currentIndex)
        }
    }
}

extension PartialJsonExtractor {
    private func extractStringValue(
        json: String,
        startIndex: String.Index
    ) -> String? {
        var result = ""
        var i = startIndex
        var escaped = false
        
        while i < json.endIndex {
            let char = json[i]
            
            if escaped {
                switch char {
                case "n":
                    result.append("\n")
                case "t":
                    result.append("\t")
                case "r":
                    result.append("\r")
                case "b":
                    result.append("\u{8}") // Backspace
                case "\"", "\\", "/":
                    result.append(char)
                case "u":
                    // Unicode escape (incomplete handling for partial JSON)
                    let remainingDistance = json.distance(from: i, to: json.endIndex)
                    if remainingDistance >= 5 {
                        let unicodeStart = json.index(after: i)
                        let unicodeEnd = json.index(unicodeStart, offsetBy: 4)
                        let unicodeHex = String(json[unicodeStart..<unicodeEnd])
                        if let unicodeValue = Int(unicodeHex, radix: 16),
                           let unicodeScalar = UnicodeScalar(unicodeValue) {
                            result.append(Character(unicodeScalar))
                            i = json.index(before: unicodeEnd)
                        } else {
                            result.append("\\u")
                        }
                    } else {
                        // Incomplete unicode escape, keep as-is
                        result.append("\\u")
                    }
                default:
                    // Unknown escape, keep the char
                    result.append(char)
                }
                escaped = false
            } else if char == "\\" {
                escaped = true
            } else if char == "\"" {
                // End of string value found
                return result
            } else {
                result.append(char)
            }
            
            i = json.index(after: i)
        }
        
        // String is incomplete (no closing quote found)
        // Return what we have so far
        return result
    }
    
    private func extractNonStringValue(
        json: String,
        startIndex: String.Index
    ) -> String? {
        var result = ""
        var i = startIndex
        
        while i < json.endIndex {
            let char = json[i]
            if char.isWhitespace || char == "," || char == "}" || char == "]" {
                // End of non-string value
                return result.isEmpty ? nil : result
            } else {
                result.append(char)
            }
            i = json.index(after: i)
        }
        
        // Value extends to end of string
        return result.isEmpty ? nil : result
    }
}