import Foundation
import AppKit
import CoreImage
import Combine

@MainActor
class ScreenCaptureViewModel: ObservableObject {
    @Published var capturedImage: NSImage?
    @Published var croppedBoardImage: NSImage?
    @Published var captureError: CaptureError?
    @Published var isCapturing = false
    @Published var isBoardDetected = false
    @Published var isAutoCapturing = true  // Default to auto-capture mode
    @Published var snapshotTaken = false   // Track if snapshot was taken
    @Published var latestSnapshot: NSImage?  // Make snapshot accessible to view
    @Published var currentPosition: ChessPosition?  // Current chess position
    @Published var recognitionConfidence: Double = 0.0  // Recognition confidence
    @Published var isAnalyzing = false  // Track analysis state
    
    private let captureManager = ScreenCapture()  // For actual capture
    private let monitorManager = ScreenCapture()  // For monitoring
    private let pieceRecognizer: PieceRecognizer
    private let boardDetector: BoardDetector
    private var captureTask: Task<Void, Never>?
    private var monitorTask: Task<Void, Never>?
    // Share CIContext to avoid creating too many Metal command queues
    private static let shared = CIContext()
    private var context: CIContext { ScreenCaptureViewModel.shared }
    
    private var detectedBoardRect: CGRect?
    private var isMonitoring = false
    
    enum CaptureError: LocalizedError {
        case chessWindowNotFound
        case boardDetectionFailed
        case noBoardDetected
        case snapshotFailed
        case recognitionFailed
        case invalidPosition
        
        var errorDescription: String? {
            switch self {
            case .chessWindowNotFound:
                return "Chess window not found"
            case .boardDetectionFailed:
                return "Failed to detect chess board"
            case .noBoardDetected:
                return "No chess board detected"
            case .snapshotFailed:
                return "Failed to take snapshot"
            case .recognitionFailed:
                return "Failed to recognize pieces"
            case .invalidPosition:
                return "Invalid chess position detected"
            }
        }
    }
    
    init() {
        // Initialize piece recognizer
        do {
            pieceRecognizer = try PieceRecognizer()
            boardDetector = BoardDetector(pieceRecognizer: pieceRecognizer)
        } catch {
            fatalError("Failed to initialize piece recognizer: \(error)")
        }
    }
    
    func takeSnapshot() async {
        print("=== SCAN BUTTON PRESSED ===")
        do {
            try await captureManager.stopCapture()
            try await captureManager.startCapture(excludeCursor: true)
            try await Task.sleep(nanoseconds: 100_000_000) // 0.1 seconds
            
            if let currentImage = captureManager.getCurrentImage() {
                print("Captured image: \(currentImage.size)")
                try await processImage(currentImage, analyzePosition: true)
                
                if let boardImage = croppedBoardImage {
                    latestSnapshot = boardImage
                    snapshotTaken = true
                    
                    if let position = currentPosition {
                        print("Successfully detected position")
                    } else {
                        print("ERROR: Failed to detect position")
                        captureError = .recognitionFailed
                    }
                } else {
                    print("ERROR: Failed to crop board image")
                    captureError = .snapshotFailed
                    currentPosition = nil
                }
            } else {
                print("ERROR: Failed to capture image")
                captureError = .snapshotFailed
                currentPosition = nil
            }
            
            try await captureManager.startCapture(excludeCursor: false)
        } catch {
            print("ERROR: Snapshot failed - \(error)")
            captureError = .snapshotFailed
            currentPosition = nil
            try? await captureManager.startCapture(excludeCursor: false)
        }
    }
    
    func startMonitoring() {
        guard !isMonitoring else { return }
        isMonitoring = true
        
        monitorTask = Task {
            monitorLoop: while !Task.isCancelled {
                do {
                    // Start monitoring capture
                    try await monitorManager.startCapture()
                    
                    while !Task.isCancelled {
                        if let image = monitorManager.getCurrentImage(),
                           let cgImage = image.cgImage(forProposedRect: nil, context: nil, hints: nil) {
                            // Check for board
                            let ciImage = CIImage(cgImage: cgImage)
                            if boardDetector.detectBoard(in: ciImage) != nil {
                                // Board detected, start capture if not already capturing
                                if !isCapturing {
                                    await startCapture()
                                }
                            }
                        }
                        try await Task.sleep(nanoseconds: 500_000_000) // 0.5 seconds
                    }
                } catch {
                    print("Monitor error: \(error)")
                    // If there's an error, wait briefly and try again
                    if !Task.isCancelled {
                        try? await Task.sleep(nanoseconds: 1_000_000_000) // 1 second
                        continue monitorLoop
                    }
                }
            }
            
            isMonitoring = false
            try? await monitorManager.stopCapture()
        }
    }
    
    func stopMonitoring() {
        monitorTask?.cancel()
        monitorTask = nil
        isMonitoring = false
        stopCapture()
        
        Task {
            try? await monitorManager.stopCapture()
        }
    }
    
    func startCapture() async {
        guard !isCapturing else { return }
        
        isCapturing = true
        captureError = nil
        snapshotTaken = false
        latestSnapshot = nil
        currentPosition = nil
        recognitionConfidence = 0.0
        
        do {
            try await captureManager.startCapture()
        } catch {
            handleCaptureError(error)
            isCapturing = false
            return
        }
        
        captureTask = Task {
            captureLoop: while !Task.isCancelled {
                do {
                    if let image = captureManager.getCurrentImage() {
                        try await processImage(image, analyzePosition: false)
                    }
                    try await Task.sleep(nanoseconds: 100_000_000) // 0.1 seconds
                } catch is CancellationError {
                    break captureLoop
                } catch {
                    // Just update the error state but continue capturing
                    if !Task.isCancelled {
                        handleCaptureError(error)
                        if isAutoCapturing && error as? CaptureError == .noBoardDetected {
                            // In auto-capture mode, stop capture but keep monitoring
                            stopCapture()
                            break captureLoop
                        }
                    }
                }
            }
            
            if isCapturing {
                isCapturing = false
            }
        }
    }
    
    func stopCapture() {
        captureTask?.cancel()
        captureTask = nil
        isCapturing = false
        captureError = nil
        isBoardDetected = false
        snapshotTaken = false
        latestSnapshot = nil
        currentPosition = nil
        recognitionConfidence = 0.0
        
        // Clean up capture session
        Task {
            try? await captureManager.stopCapture()
        }
    }
    
    private func processImage(_ image: NSImage, analyzePosition: Bool = false) async throws {
        guard let cgImage = image.cgImage(forProposedRect: nil, context: nil, hints: nil) else {
            print("ERROR: Failed to convert image for processing")
            throw CaptureError.boardDetectionFailed
        }
        
        let ciImage = CIImage(cgImage: cgImage)
        
        if let boardRect = boardDetector.detectBoard(in: ciImage) {
            isBoardDetected = true
            self.detectedBoardRect = boardRect
            
            let croppedImage = ciImage.cropped(to: boardRect)
            updateImages(ciImage: ciImage, croppedImage: croppedImage)
            
            if analyzePosition && !isAnalyzing {
                isAnalyzing = true
                do {
                    print("Analyzing board position...")
                    let position = try await boardDetector.analyzeBoard(in: ciImage)
                    currentPosition = position
                    recognitionConfidence = 1.0
                    isAnalyzing = false
                } catch {
                    print("ERROR: Position analysis failed - \(error)")
                    isAnalyzing = false
                    currentPosition = nil
                    recognitionConfidence = 0.0
                    if error is BoardDetectionError {
                        throw CaptureError.recognitionFailed
                    } else {
                        throw error
                    }
                }
            }
        } else {
            isBoardDetected = false
            self.detectedBoardRect = nil
            currentPosition = nil
            recognitionConfidence = 0.0
            
            if let cgImage = context.createCGImage(ciImage, from: ciImage.extent) {
                self.capturedImage = NSImage(cgImage: cgImage, size: .zero)
            }
            self.croppedBoardImage = nil
            
            if isAutoCapturing {
                throw CaptureError.noBoardDetected
            }
        }
    }
    
    private func updateImages(ciImage: CIImage, croppedImage: CIImage) {
        if let cgImage = context.createCGImage(ciImage, from: ciImage.extent) {
            let capturedNSImage = NSImage(cgImage: cgImage, size: .zero)
            self.capturedImage = capturedNSImage
        }
        
        if let cgCroppedImage = context.createCGImage(croppedImage, from: croppedImage.extent) {
            let croppedNSImage = NSImage(cgImage: cgCroppedImage, size: .zero)
            self.croppedBoardImage = croppedNSImage
        }
    }
    
    private func handleCaptureError(_ error: Error) {
        if let captureError = error as? CaptureError {
            self.captureError = captureError
        } else {
            self.captureError = .boardDetectionFailed
        }
    }
}
