//
//  MaskingService.swift
//  CheapRetouch
//
//  Wrapper around Vision framework for generating masks from user taps.
//

import Foundation
import Vision
import CoreImage
import UIKit

actor MaskingService {

    enum MaskingError: Error, LocalizedError {
        case imageConversionFailed
        case requestFailed(Error)
        case noMaskGenerated
        case instanceNotFound

        var errorDescription: String? {
            switch self {
            case .imageConversionFailed:
                return "Failed to convert image for processing"
            case .requestFailed(let error):
                return "Vision request failed: \(error.localizedDescription)"
            case .noMaskGenerated:
                return "No mask was generated for the selected area"
            case .instanceNotFound:
                return "No object found at the tap location"
            }
        }
    }

    func generatePersonMask(at point: CGPoint, in image: CGImage) async throws -> CGImage? {
        try await generateForegroundMask(at: point, in: image)
    }

    struct MaskResult {
        let mask: CGImage?
        let instanceCount: Int
        let confidence: Float // 0.0 to 1.0, based on mask coverage relative to image
    }

    func generatePersonMaskWithCount(at point: CGPoint, in image: CGImage) async throws -> MaskResult {
        DebugLogger.processing("MaskingService.generatePersonMaskWithCount at \(point)")
        DebugLogger.log("Image size: \(image.width)x\(image.height)")

        let request = VNGenerateForegroundInstanceMaskRequest()

        let handler = VNImageRequestHandler(cgImage: image, options: [:])

        do {
            DebugLogger.log("Performing Vision request...")
            try handler.perform([request])
            DebugLogger.log("Vision request completed")
        } catch {
            DebugLogger.error("Vision request failed", error: error)
            throw MaskingError.requestFailed(error)
        }

        guard let result = request.results?.first else {
            return MaskResult(mask: nil, instanceCount: 0, confidence: 0)
        }

        let allInstances = result.allInstances
        let instanceCount = allInstances.count

        let visionPoint = CGPoint(
            x: point.x / CGFloat(image.width),
            y: 1.0 - point.y / CGFloat(image.height)
        )

        // Find instance at tap point
        var targetInstance: IndexSet?

        for instance in allInstances {
            let indexSet = IndexSet(integer: instance)
            if let maskPixelBuffer = try? result.generateScaledMaskForImage(forInstances: indexSet, from: handler) {
                if isPoint(visionPoint, inMask: maskPixelBuffer, imageSize: CGSize(width: image.width, height: image.height)) {
                    targetInstance = indexSet
                    break
                }
            }
        }

        guard let instance = targetInstance else {
            return MaskResult(mask: nil, instanceCount: instanceCount, confidence: 0)
        }

        let maskPixelBuffer = try result.generateScaledMaskForImage(forInstances: instance, from: handler)
        let maskImage = convertPixelBufferToCGImage(maskPixelBuffer)

        // Calculate confidence based on mask quality (coverage ratio)
        let confidence = calculateMaskConfidence(maskPixelBuffer, imageSize: CGSize(width: image.width, height: image.height))

        return MaskResult(mask: maskImage, instanceCount: instanceCount, confidence: confidence)
    }

    private func calculateMaskConfidence(_ pixelBuffer: CVPixelBuffer, imageSize: CGSize) -> Float {
        CVPixelBufferLockBaseAddress(pixelBuffer, .readOnly)
        defer { CVPixelBufferUnlockBaseAddress(pixelBuffer, .readOnly) }

        let width = CVPixelBufferGetWidth(pixelBuffer)
        let height = CVPixelBufferGetHeight(pixelBuffer)

        guard let baseAddress = CVPixelBufferGetBaseAddress(pixelBuffer) else {
            return 0.5 // Default medium confidence if can't read
        }

        let bytesPerRow = CVPixelBufferGetBytesPerRow(pixelBuffer)
        var maskPixelCount = 0
        var edgePixelCount = 0

        // Count mask pixels and check edge sharpness
        for y in 0..<height {
            for x in 0..<width {
                let pixelOffset = y * bytesPerRow + x
                let pixelValue = baseAddress.load(fromByteOffset: pixelOffset, as: UInt8.self)

                if pixelValue > 127 {
                    maskPixelCount += 1
                }

                // Check for edge pixels (values between 50-200 indicate soft edges)
                if pixelValue > 50 && pixelValue < 200 {
                    edgePixelCount += 1
                }
            }
        }

        let totalPixels = width * height
        let maskRatio = Float(maskPixelCount) / Float(totalPixels)
        let edgeRatio = Float(edgePixelCount) / Float(max(1, maskPixelCount))

        // Confidence is higher when:
        // - Mask covers reasonable portion (not too small, not too large)
        // - Edge pixels are minimal (sharp edges = confident detection)
        var confidence: Float = 1.0

        // Penalize very small masks (< 1% of image)
        if maskRatio < 0.01 {
            confidence -= 0.3
        }

        // Penalize very large masks (> 50% of image)
        if maskRatio > 0.5 {
            confidence -= 0.2
        }

        // Penalize fuzzy edges
        if edgeRatio > 0.3 {
            confidence -= 0.2
        }

        return max(0.0, min(1.0, confidence))
    }

    func generateForegroundMask(at point: CGPoint, in image: CGImage) async throws -> CGImage? {
        DebugLogger.processing("generateForegroundMask at \(point)")
        DebugLogger.log("Image: \(image.width)x\(image.height)")

        let request = VNGenerateForegroundInstanceMaskRequest()
        DebugLogger.log("Created VNGenerateForegroundInstanceMaskRequest")

        let handler = VNImageRequestHandler(cgImage: image, options: [:])
        DebugLogger.log("Created VNImageRequestHandler, performing request...")

        do {
            try handler.perform([request])
            DebugLogger.log("Vision request completed successfully")
        } catch {
            DebugLogger.error("Vision request failed", error: error)
            throw MaskingError.requestFailed(error)
        }

        guard let result = request.results?.first else {
            DebugLogger.log("No results from Vision request")
            return nil
        }

        DebugLogger.log("Got Vision result with \(result.allInstances.count) instances")

        // Normalize point to Vision coordinates (0-1, origin bottom-left)
        let visionPoint = CGPoint(
            x: point.x / CGFloat(image.width),
            y: 1.0 - point.y / CGFloat(image.height)
        )
        DebugLogger.log("Vision point: \(visionPoint)")

        // Find instance at tap point
        let instances = result.allInstances
        var targetInstance: IndexSet?

        DebugLogger.log("Checking \(instances.count) instances for tap point...")
        for instance in instances {
            let indexSet = IndexSet(integer: instance)
            if let maskPixelBuffer = try? result.generateScaledMaskForImage(forInstances: indexSet, from: handler) {
                // Check if point is within this instance's mask
                if isPoint(visionPoint, inMask: maskPixelBuffer, imageSize: CGSize(width: image.width, height: image.height)) {
                    DebugLogger.log("Found instance \(instance) at tap point")
                    targetInstance = indexSet
                    break
                }
            }
        }

        guard let instance = targetInstance else {
            DebugLogger.log("No instance found at tap point")
            return nil
        }

        DebugLogger.log("Generating mask for selected instance...")
        let maskPixelBuffer = try result.generateScaledMaskForImage(forInstances: instance, from: handler)
        DebugLogger.log("Mask generated, converting to CGImage...")
        let cgImage = convertPixelBufferToCGImage(maskPixelBuffer)
        DebugLogger.imageInfo("Generated mask", image: cgImage)
        return cgImage
    }

    func generateAllForegroundMasks(in image: CGImage) async throws -> CGImage? {
        let request = VNGenerateForegroundInstanceMaskRequest()

        let handler = VNImageRequestHandler(cgImage: image, options: [:])

        do {
            try handler.perform([request])
        } catch {
            throw MaskingError.requestFailed(error)
        }

        guard let result = request.results?.first else {
            return nil
        }

        let allInstances = result.allInstances
        guard !allInstances.isEmpty else {
            return nil
        }

        let maskPixelBuffer = try result.generateScaledMaskForImage(forInstances: IndexSet(allInstances), from: handler)
        return convertPixelBufferToCGImage(maskPixelBuffer)
    }

    private func isPoint(_ point: CGPoint, inMask pixelBuffer: CVPixelBuffer, imageSize: CGSize) -> Bool {
        CVPixelBufferLockBaseAddress(pixelBuffer, .readOnly)
        defer { CVPixelBufferUnlockBaseAddress(pixelBuffer, .readOnly) }

        let maskWidth = CVPixelBufferGetWidth(pixelBuffer)
        let maskHeight = CVPixelBufferGetHeight(pixelBuffer)
        let bytesPerRow = CVPixelBufferGetBytesPerRow(pixelBuffer)
        let bytesPerPixel = bytesPerRow / maskWidth

        DebugLogger.log("isPoint check: visionPoint=\(point), maskSize=\(maskWidth)x\(maskHeight), bytesPerPixel=\(bytesPerPixel)")

        // point is in Vision coordinates (0-1 normalized, origin bottom-left)
        // Mask pixel buffer has origin at top-left
        // Scale point to mask dimensions and flip Y from Vision coords to image/mask coords
        let x = Int(point.x * CGFloat(maskWidth))
        let y = Int((1.0 - point.y) * CGFloat(maskHeight))

        DebugLogger.log("Pixel coords: x=\(x), y=\(y) (from visionY=\(point.y) -> imageY=\(1.0 - point.y))")

        guard x >= 0, x < maskWidth, y >= 0, y < maskHeight else {
            DebugLogger.log("Pixel coords out of bounds")
            return false
        }

        guard let baseAddress = CVPixelBufferGetBaseAddress(pixelBuffer) else {
            DebugLogger.log("Failed to get pixel buffer base address")
            return false
        }
        
        // Calculate offset accounting for bytes per pixel
        let pixelOffset = y * bytesPerRow + x * bytesPerPixel
        let pixelValue = baseAddress.load(fromByteOffset: pixelOffset, as: UInt8.self)
        
        // Also sample nearby pixels to see if we're close to the mask
        var nearbyValues: [UInt8] = []
        for dy in -2...2 {
            for dx in -2...2 {
                let nx = x + dx
                let ny = y + dy
                if nx >= 0, nx < maskWidth, ny >= 0, ny < maskHeight {
                    let offset = ny * bytesPerRow + nx * bytesPerPixel
                    nearbyValues.append(baseAddress.load(fromByteOffset: offset, as: UInt8.self))
                }
            }
        }
        let maxNearby = nearbyValues.max() ?? 0
        
        DebugLogger.log("Pixel value at (\(x), \(y)): \(pixelValue), maxNearby=\(maxNearby) (threshold 127)")
        return pixelValue > 127
    }

    private func convertPixelBufferToCGImage(_ pixelBuffer: CVPixelBuffer) -> CGImage? {
        let ciImage = CIImage(cvPixelBuffer: pixelBuffer)
        let context = CIContext()
        return context.createCGImage(ciImage, from: ciImage.extent)
    }
}
