import Foundation
import CoreGraphics
import Vision

public struct JointResult: Codable, Sendable {
    public var name: String
    public var point: Point2
    public var position3D: [Double]?
    public var confidence2D: Double?
    public var parent: String?
}
public struct PoseResult: Codable, Sendable {
    public var kind: String
    public var status: String
    public var requestRevision: Int
    public var durationSeconds: Double
    public var joints: [JointResult]
    public var message: String?
    public var bodyHeight: Double?
    public var heightEstimation: String?
    public var candidateCount: Int
    public var cameraOrigin: [Double]?
}
public enum VisionPoseEstimator {
    public static func estimate3D(image: CGImage, crop: PixelRect) -> PoseResult {
        let request = VNDetectHumanBodyPose3DRequest()
        request.revision = VNDetectHumanBodyPose3DRequestRevision1
        let start = Date()
        do {
            try VNImageRequestHandler(cgImage: image, orientation: .up).perform([request])
            guard let observation = request.results?.first else {
                return result(kind: "3D", request: request, start: start, status: "noPose")
            }
            let names: [(String, VNHumanBodyPose3DObservation.JointName)] = [
                ("root", .root), ("spine", .spine), ("center_shoulder", .centerShoulder),
                ("center_head", .centerHead), ("top_head", .topHead),
                ("left_shoulder", .leftShoulder), ("left_elbow", .leftElbow), ("left_wrist", .leftWrist),
                ("right_shoulder", .rightShoulder), ("right_elbow", .rightElbow), ("right_wrist", .rightWrist),
                ("left_hip", .leftHip), ("left_knee", .leftKnee), ("left_ankle", .leftAnkle),
                ("right_hip", .rightHip), ("right_knee", .rightKnee), ("right_ankle", .rightAnkle)]
            let joints = try names.map { label, name in
                let p = try observation.recognizedPoint(name)
                let projected = try observation.pointInImage(name)
                let t = p.position.columns.3
                return JointResult(name: label, point: crop.sourcePoint(normalizedX: projected.x, normalizedY: projected.y),
                                   position3D: [Double(t.x), Double(t.y), Double(t.z)], confidence2D: nil,
                                   parent: names.first(where: { $0.1 == p.parentJoint })?.0)
            }
            var value = result(kind: "3D", request: request, start: start, status: "needsReview")
            value.joints = joints
            value.bodyHeight = Double(observation.bodyHeight)
            value.heightEstimation = observation.heightEstimation == .measured ? "measured" : "reference"
            value.candidateCount = request.results?.count ?? 0
            let m = observation.cameraOriginMatrix
            value.cameraOrigin = (0..<4).flatMap { c in (0..<4).map { r in Double(m[c][r]) } }
            return value
        } catch { return result(kind: "3D", request: request, start: start, status: "error", message: error.localizedDescription) }
    }
    public static func estimate2D(image: CGImage, crop: PixelRect) -> PoseResult {
        let request = VNDetectHumanBodyPoseRequest()
        request.revision = VNDetectHumanBodyPoseRequestRevision1
        let start = Date()
        do {
            try VNImageRequestHandler(cgImage: image, orientation: .up).perform([request])
            let observations = request.results ?? []
            guard let observation = observations.first else { return result(kind: "2D", request: request, start: start, status: "noPose") }
            let names: [(String, VNHumanBodyPoseObservation.JointName)] = [
                ("root", .root), ("neck", .neck), ("nose", .nose),
                ("left_shoulder", .leftShoulder), ("left_elbow", .leftElbow), ("left_wrist", .leftWrist),
                ("right_shoulder", .rightShoulder), ("right_elbow", .rightElbow), ("right_wrist", .rightWrist),
                ("left_hip", .leftHip), ("left_knee", .leftKnee), ("left_ankle", .leftAnkle),
                ("right_hip", .rightHip), ("right_knee", .rightKnee), ("right_ankle", .rightAnkle)]
            var value = result(kind: "2D", request: request, start: start, status: "needsReview")
            value.candidateCount = observations.count
            // Multiple candidates require explicit selection; never silently pick a different person.
            if observations.count != 1 {
                value.status = "ambiguous"
                value.message = "Multiple 2D candidates; tighten the selected rider crop."
                return value
            }
            value.joints = try names.compactMap { label, name in
                let p = try observation.recognizedPoint(name)
                guard p.confidence > 0 else { return nil }
                return JointResult(name: label, point: crop.sourcePoint(normalizedX: p.x, normalizedY: p.y),
                                   position3D: nil, confidence2D: Double(p.confidence), parent: nil)
            }
            return value
        } catch { return result(kind: "2D", request: request, start: start, status: "error", message: error.localizedDescription) }
    }
    private static func result(kind: String, request: VNRequest, start: Date, status: String, message: String? = nil) -> PoseResult {
        PoseResult(kind: kind, status: status, requestRevision: request.revision, durationSeconds: Date().timeIntervalSince(start),
                   joints: [], message: message, bodyHeight: nil, heightEstimation: nil, candidateCount: 0, cameraOrigin: nil)
    }
}