//
//  ContourService.swift
//  CheapRetouch
//
//  Service for detecting and scoring wire/line contours.
//

import Foundation
import Vision
import CoreGraphics
import UIKit

actor ContourService {

    struct ScoredContour {
        let contour: VNContour
        let score: Float
    }

    private let proximityWeight: Float = 0.3
    private let aspectWeight: Float = 0.3
    private let straightnessWeight: Float = 0.2
    private let lengthWeight: Float = 0.2

    private let minimumScore: Float = 0.3

    func detectContours(in image: CGImage) async throws -> [VNContour] {
        let request = VNDetectContoursRequest()
        request.contrastAdjustment = 1.0
        request.detectsDarkOnLight = true

        let handler = VNImageRequestHandler(cgImage: image, options: [:])
        try handler.perform([request])

        guard let results = request.results?.first else {
            return []
        }

        return collectAllContours(from: results)
    }

    func findBestWireContour(at point: CGPoint, from contours: [VNContour], imageSize: CGSize) -> VNContour? {
        let normalizedPoint = CGPoint(
            x: point.x / imageSize.width,
            y: 1.0 - point.y / imageSize.height
        )

        let scoredContours = contours.compactMap { contour -> ScoredContour? in
            let score = scoreContour(contour, relativeTo: normalizedPoint)
            guard score >= minimumScore else { return nil }
            return ScoredContour(contour: contour, score: score)
        }

        return scoredContours.max(by: { $0.score < $1.score })?.contour
    }

    func scoreContour(_ contour: VNContour, relativeTo point: CGPoint) -> Float {
        let points = contour.normalizedPoints

        guard points.count >= 2 else { return 0 }

        let proximityScore = calculateProximityScore(points: points, to: point)
        let aspectScore = calculateAspectScore(points: points)
        let straightnessScore = calculateStraightnessScore(points: points)
        let lengthScore = calculateLengthScore(points: points)

        return proximityScore * proximityWeight +
               aspectScore * aspectWeight +
               straightnessScore * straightnessWeight +
               lengthScore * lengthWeight
    }

    func contourToMask(_ contour: VNContour, width: Int, imageSize: CGSize) -> CGImage? {
        let maskWidth = Int(imageSize.width)
        let maskHeight = Int(imageSize.height)

        guard let context = CGContext(
            data: nil,
            width: maskWidth,
            height: maskHeight,
            bitsPerComponent: 8,
            bytesPerRow: maskWidth,
            space: CGColorSpaceCreateDeviceGray(),
            bitmapInfo: CGImageAlphaInfo.none.rawValue
        ) else {
            return nil
        }

        context.setFillColor(gray: 0, alpha: 1)
        context.fill(CGRect(x: 0, y: 0, width: maskWidth, height: maskHeight))

        context.setStrokeColor(gray: 1, alpha: 1)
        context.setLineWidth(CGFloat(width))
        context.setLineCap(.round)
        context.setLineJoin(.round)

        let points = contour.normalizedPoints
        guard let firstPoint = points.first else { return nil }

        context.beginPath()
        context.move(to: CGPoint(
            x: CGFloat(firstPoint.x) * imageSize.width,
            y: CGFloat(firstPoint.y) * imageSize.height
        ))

        for point in points.dropFirst() {
            context.addLine(to: CGPoint(
                x: CGFloat(point.x) * imageSize.width,
                y: CGFloat(point.y) * imageSize.height
            ))
        }

        context.strokePath()

        return context.makeImage()
    }

    // MARK: - Private Scoring Methods

    private func calculateProximityScore(points: [SIMD2<Float>], to target: CGPoint) -> Float {
        var minDistance: Float = .greatestFiniteMagnitude

        for point in points {
            let dx = Float(target.x) - point.x
            let dy = Float(target.y) - point.y
            let distance = sqrt(dx * dx + dy * dy)
            minDistance = min(minDistance, distance)
        }

        // Score decreases with distance, max at 0, 0 at distance > 0.2
        return max(0, 1.0 - minDistance * 5)
    }

    private func calculateAspectScore(points: [SIMD2<Float>]) -> Float {
        guard points.count >= 2 else { return 0 }

        var minX: Float = .greatestFiniteMagnitude
        var maxX: Float = -.greatestFiniteMagnitude
        var minY: Float = .greatestFiniteMagnitude
        var maxY: Float = -.greatestFiniteMagnitude

        for point in points {
            minX = min(minX, point.x)
            maxX = max(maxX, point.x)
            minY = min(minY, point.y)
            maxY = max(maxY, point.y)
        }

        let width = maxX - minX
        let height = maxY - minY

        let length = calculatePathLength(points: points)
        let boundingDiagonal = sqrt(width * width + height * height)

        guard boundingDiagonal > 0 else { return 0 }

        // Wire-like shapes have length close to their bounding diagonal
        // and small perpendicular extent
        let minDimension = min(width, height)
        let maxDimension = max(width, height)

        guard maxDimension > 0 else { return 0 }

        let aspectRatio = minDimension / maxDimension
        // Low aspect ratio (thin) scores high
        return max(0, 1.0 - aspectRatio * 2)
    }

    private func calculateStraightnessScore(points: [SIMD2<Float>]) -> Float {
        guard points.count >= 3 else { return 1.0 }

        var totalAngleChange: Float = 0

        for i in 1..<(points.count - 1) {
            let prev = points[i - 1]
            let curr = points[i]
            let next = points[i + 1]

            let v1 = SIMD2<Float>(curr.x - prev.x, curr.y - prev.y)
            let v2 = SIMD2<Float>(next.x - curr.x, next.y - curr.y)

            let len1 = sqrt(v1.x * v1.x + v1.y * v1.y)
            let len2 = sqrt(v2.x * v2.x + v2.y * v2.y)

            guard len1 > 0, len2 > 0 else { continue }

            let dot = (v1.x * v2.x + v1.y * v2.y) / (len1 * len2)
            let angle = acos(min(1, max(-1, dot)))
            totalAngleChange += angle
        }

        let averageAngleChange = totalAngleChange / Float(points.count - 2)
        // Straight lines have low angle change
        return max(0, 1.0 - averageAngleChange / .pi)
    }

    private func calculateLengthScore(points: [SIMD2<Float>]) -> Float {
        let length = calculatePathLength(points: points)
        // Longer contours score higher, normalized to typical wire length
        return min(1.0, length * 2)
    }

    private func calculatePathLength(points: [SIMD2<Float>]) -> Float {
        var length: Float = 0

        for i in 1..<points.count {
            let dx = points[i].x - points[i - 1].x
            let dy = points[i].y - points[i - 1].y
            length += sqrt(dx * dx + dy * dy)
        }

        return length
    }

    private func collectAllContours(from observation: VNContoursObservation) -> [VNContour] {
        var contours: [VNContour] = []

        func collect(_ contour: VNContour) {
            contours.append(contour)
            for i in 0..<contour.childContourCount {
                if let child = try? contour.childContour(at: i) {
                    collect(child)
                }
            }
        }

        for i in 0..<observation.contourCount {
            if let contour = try? observation.contour(at: i) {
                collect(contour)
            }
        }

        return contours
    }
}
