import Cocoa
import ApplicationServices

public struct MacInputState: Codable {
    public let source: String
    public let app: String
    public let text: String
    public let cursor: Int
    public let selection: Int
    public let revision: Int64
    public let timestamp: Double
    
    public init(app: String, text: String, cursor: Int, selection: Int, revision: Int64) {
        self.source = "mac"
        self.app = app
        self.text = text
        self.cursor = cursor
        self.selection = selection
        self.revision = revision
        self.timestamp = Date().timeIntervalSince1970
    }
}

public final class FocusedInputSync {
    public static let shared = FocusedInputSync()
    
    private let systemWideElement: AXUIElement
    private var isApplyingRemoteChange: Bool = false
    private var lastObservedHash: Int = 0
    private var localRevision: Int64 = 0
    
    private init() {
        self.systemWideElement = AXUIElementCreateSystemWide()
        enableGlobalAccessibility()
    }
    
    private func enableGlobalAccessibility() {
        AXUIElementSetAttributeValue(systemWideElement, "AXEnhancedUserInterface" as CFString, kCFBooleanTrue)
        AXUIElementSetAttributeValue(systemWideElement, "AXManualAccessibility" as CFString, kCFBooleanTrue)
    }
    
    private func findFocusedDescendant(_ elem: AXUIElement) -> AXUIElement? {
        var isFocusedObj: CFTypeRef?
        if AXUIElementCopyAttributeValue(elem, kAXFocusedAttribute as CFString, &isFocusedObj) == .success,
           let isFocused = isFocusedObj as? Bool, isFocused {
            var roleObj: CFTypeRef?
            AXUIElementCopyAttributeValue(elem, kAXRoleAttribute as CFString, &roleObj)
            let role = (roleObj as? String) ?? ""
            if role != "AXWindow" && role != "AXApplication" && role != "AXGroup" && role != "AXSplitGroup" && role != "AXScrollArea" {
                return elem
            }
        }
        
        var childrenObj: CFTypeRef?
        if AXUIElementCopyAttributeValue(elem, kAXChildrenAttribute as CFString, &childrenObj) == .success,
           let children = childrenObj as? [AXUIElement] {
            for child in children {
                if let found = findFocusedDescendant(child) {
                    return found
                }
            }
        }
        return nil
    }
    
    public func getFocusedElement() -> (AXUIElement, String)? {
        guard let frontApp = NSWorkspace.shared.frontmostApplication else { return nil }
        let appName = frontApp.localizedName ?? "App"
        let appElem = AXUIElementCreateApplication(frontApp.processIdentifier)
        
        AXUIElementSetAttributeValue(appElem, "AXEnhancedUserInterface" as CFString, kCFBooleanTrue)
        AXUIElementSetAttributeValue(appElem, "AXManualAccessibility" as CFString, kCFBooleanTrue)
        
        var targetElem: AXUIElement?
        
        // 1. Try system wide focused element
        var focusedObj: CFTypeRef?
        if AXUIElementCopyAttributeValue(systemWideElement, kAXFocusedUIElementAttribute as CFString, &focusedObj) == .success,
           let obj = focusedObj {
            let elem = obj as! AXUIElement
            var roleObj: CFTypeRef?
            AXUIElementCopyAttributeValue(elem, kAXRoleAttribute as CFString, &roleObj)
            let role = (roleObj as? String) ?? ""
            if role != "AXWindow" && role != "AXApplication" {
                targetElem = elem
            }
        }
        
        // 2. Try App focused element
        if targetElem == nil {
            if AXUIElementCopyAttributeValue(appElem, kAXFocusedUIElementAttribute as CFString, &focusedObj) == .success,
               let obj = focusedObj {
                let elem = obj as! AXUIElement
                var roleObj: CFTypeRef?
                AXUIElementCopyAttributeValue(elem, kAXRoleAttribute as CFString, &roleObj)
                let role = (roleObj as? String) ?? ""
                if role != "AXWindow" && role != "AXApplication" {
                    targetElem = elem
                }
            }
        }
        
        // 3. Recursive search in tree
        if targetElem == nil {
            targetElem = findFocusedDescendant(appElem)
        }
        
        if let elem = targetElem {
            return (elem, appName)
        }
        return nil
    }
    
    /// Inspects the current focused element and returns a state snapshot if changed
    public func inspectCurrentState() -> MacInputState? {
        guard !isApplyingRemoteChange else { return nil }
        guard let (elem, appName) = getFocusedElement() else { return nil }
        
        // Extract text
        var text = ""
        var valObj: CFTypeRef?
        if AXUIElementCopyAttributeValue(elem, kAXValueAttribute as CFString, &valObj) == .success,
           let val = valObj {
            if let str = val as? String {
                text = str
            } else if let attrStr = val as? NSAttributedString {
                text = attrStr.string
            }
        }
        
        if text.isEmpty {
            var countObj: CFTypeRef?
            if AXUIElementCopyAttributeValue(elem, kAXNumberOfCharactersAttribute as CFString, &countObj) == .success,
               let count = countObj as? Int, count > 0 {
                var range = CFRange(location: 0, length: count)
                if let axRange = AXValueCreate(.cfRange, &range) {
                    var strObj: CFTypeRef?
                    if AXUIElementCopyParameterizedAttributeValue(elem, kAXStringForRangeParameterizedAttribute as CFString, axRange, &strObj) == .success,
                       let str = strObj as? String {
                        text = str
                    }
                }
            }
        }
        
        // Extract cursor & selection
        var cursor = text.count
        var selLen = 0
        var rangeObj: CFTypeRef?
        if AXUIElementCopyAttributeValue(elem, kAXSelectedTextRangeAttribute as CFString, &rangeObj) == .success,
           let axRange = rangeObj {
            var range = CFRange()
            if AXValueGetValue(axRange as! AXValue, .cfRange, &range) {
                cursor = range.location
                selLen = range.length
            }
        }
        
        // Check hash to prevent echo loops
        let currentHash = "\(appName)_\(text)_\(cursor)_\(selLen)".hashValue
        guard currentHash != lastObservedHash else { return nil }
        lastObservedHash = currentHash
        localRevision += 1
        
        return MacInputState(app: appName, text: text, cursor: cursor, selection: selLen, revision: localRevision)
    }
    
    /// Applies updated full text or inserts at cursor position directly into active Mac input
    @discardableResult
    public func applyRemoteUpdate(text: String, cursor: Int? = nil, isFullReplace: Bool = true) -> Bool {
        isApplyingRemoteChange = true
        defer {
            DispatchQueue.main.asyncAfter(deadline: .now() + 0.12) {
                self.isApplyingRemoteChange = false
            }
        }
        
        // 1. Flatten all newlines and carriage returns
        var cleanText = text.components(separatedBy: .newlines).joined(separator: " ")
        cleanText = cleanText.replacingOccurrences(of: "\t", with: " ")
        cleanText = cleanText.replacingOccurrences(of: "\\s+", with: " ", options: .regularExpression)
        cleanText = cleanText.trimmingCharacters(in: .whitespacesAndNewlines)
        
        guard let (elem, appName) = getFocusedElement() else {
            // Fallback: clipboard paste
            return pasteViaKeystroke(cleanText, isFullReplace: isFullReplace)
        }
        
        if isFullReplace {
            // Try setting AXValue directly
            let setErr = AXUIElementSetAttributeValue(elem, kAXValueAttribute as CFString, cleanText as CFTypeRef)
            if setErr == .success {
                if let cursor = cursor {
                    var range = CFRange(location: min(cursor, cleanText.count), length: 0)
                    if let axRange = AXValueCreate(.cfRange, &range) {
                        AXUIElementSetAttributeValue(elem, kAXSelectedTextRangeAttribute as CFString, axRange)
                    }
                }
                lastObservedHash = "\(appName)_\(cleanText)_\(cursor ?? cleanText.count)_0".hashValue
                print("FocusedInputSync: ✅ Updated AXValue for \(appName)")
                return true
            }
        } else {
            // Try setting selected text
            let setSelErr = AXUIElementSetAttributeValue(elem, kAXSelectedTextAttribute as CFString, cleanText as CFTypeRef)
            if setSelErr == .success {
                lastObservedHash = "\(appName)_\(cleanText)_\(cursor ?? cleanText.count)_0".hashValue
                print("FocusedInputSync: ✅ Inserted text via AXSelectedText for \(appName)")
                return true
            }
        }
        
        return pasteViaKeystroke(cleanText, isFullReplace: isFullReplace)
    }
    
    private func pasteViaKeystroke(_ text: String, isFullReplace: Bool) -> Bool {
        let pasteboard = NSPasteboard.general
        pasteboard.clearContents()
        pasteboard.setString(text, forType: .string)
        
        let src = CGEventSource(stateID: .hidSystemState)
        
        if isFullReplace {
            // Select all: Cmd + A
            let aDown = CGEvent(keyboardEventSource: src, virtualKey: 0, keyDown: true)
            aDown?.flags = .maskCommand
            let aUp = CGEvent(keyboardEventSource: src, virtualKey: 0, keyDown: false)
            aDown?.post(tap: .cghidEventTap)
            aUp?.post(tap: .cghidEventTap)
            
            usleep(25000)
        }
        
        // Paste: Cmd + V
        let vDown = CGEvent(keyboardEventSource: src, virtualKey: 9, keyDown: true)
        vDown?.flags = .maskCommand
        let vUp = CGEvent(keyboardEventSource: src, virtualKey: 9, keyDown: false)
        vDown?.post(tap: .cghidEventTap)
        vUp?.post(tap: .cghidEventTap)
        
        return true
    }
}