import Foundation
import CoreGraphics
import ImageIO
import CoreText
import UniformTypeIdentifiers

public enum InspectionRenderer {
    // Fixed union of supported 2D/3D joints. Never renumber a partial detection.
    public static let jointNames = ["root", "spine", "center_shoulder", "center_head", "top_head",
        "left_shoulder", "left_elbow", "left_wrist", "right_shoulder", "right_elbow", "right_wrist",
        "left_hip", "left_knee", "left_ankle", "right_hip", "right_knee", "right_ankle", "neck", "nose"]
    public static func jointID(_ name: String) -> Int? { jointNames.firstIndex(of: name).map { $0 + 1 } }

    public static let bones = [("root","left_hip"),("root","right_hip"),("root","spine"),("spine","center_shoulder"),
        ("center_shoulder","center_head"),("center_head","top_head"),("center_shoulder","left_shoulder"),("center_shoulder","right_shoulder"),
        ("left_shoulder","right_shoulder"),("left_shoulder","left_elbow"),("left_elbow","left_wrist"),
        ("right_shoulder","right_elbow"),("right_elbow","right_wrist"),("left_hip","left_knee"),("left_knee","left_ankle"),
        ("right_hip","right_knee"),("right_knee","right_ankle"),("left_hip","right_hip")]
    public static func write(image: CGImage, panel: PixelRect, crop: PixelRect, pose: PoseResult, annotations: [Landmark], url: URL) throws {
        let height = max(image.height, 560)
        guard let context = CGContext(data: nil, width: image.width + 310, height: height, bitsPerComponent: 8, bytesPerRow: 0,
                                      space: CGColorSpaceCreateDeviceRGB(), bitmapInfo: CGImageAlphaInfo.premultipliedLast.rawValue) else {
            throw PoseError.invalidInput("Cannot allocate renderer.")
        }
        context.setFillColor(CGColor(gray: 0.07, alpha: 1))
        context.fill(CGRect(x: 0, y: 0, width: image.width + 310, height: height))
        context.draw(image, in: CGRect(x: 0, y: height - image.height, width: image.width, height: image.height))
        context.translateBy(x: 0,y: Double(height)); context.scaleBy(x: 1,y: -1)
        context.saveGState()
        context.clip(to: CGRect(x: 0, y: 0, width: image.width, height: image.height))
        func local(_ p: Point2) -> CGPoint { CGPoint(x: p.x-panel.x,y: p.y-panel.y) }
        context.setStrokeColor(CGColor(red: 1,green: 1,blue: 0,alpha: 0.9)); context.setLineWidth(1)
        context.stroke(CGRect(x: crop.x-panel.x,y: crop.y-panel.y,width: crop.width,height: crop.height))
        let points = Dictionary(uniqueKeysWithValues: pose.joints.filter { ($0.confidence2D ?? 1) >= 0.2 }.map { ($0.name,$0.point) })
        context.setStrokeColor(CGColor(red: 0,green: 1,blue: 1,alpha: 1)); context.setLineWidth(1.5)
        for (a,b) in bones {
            if let p = points[a], let q = points[b] { context.move(to: local(p)); context.addLine(to: local(q)); context.strokePath() }
        }
        for joint in pose.joints {
            guard points[joint.name] != nil else { continue }
            let point = joint.point
            let p=local(point)
            context.setFillColor(CGColor(red: 0,green: 1,blue: 1,alpha: 1))
            context.fillEllipse(in: CGRect(x: p.x-2,y: p.y-2,width: 4,height: 4))
            if let id = jointID(joint.name) {
                label(String(id), at: CGPoint(x: min(max(p.x+3, 2), Double(image.width)-14),
                                                y: min(max(p.y-3, 9), Double(image.height)-2)), context: context)
            }
        }
        context.setStrokeColor(CGColor(red: 1,green: 0.2,blue: 0.8,alpha: 1)); context.setLineWidth(1)
        for mark in annotations {
            let p=local(mark.point), r=max(mark.uncertaintyPixels,2)
            context.strokeEllipse(in: CGRect(x: p.x-r,y: p.y-r,width: r*2,height: r*2))
        }
        context.restoreGState()
        let x = Double(image.width) + 16
        label("\(pose.kind) · UNREVIEWED", at: CGPoint(x: x, y: 25), context: context, size: 14)
        label("Status: \(pose.status)", at: CGPoint(x: x, y: 47), context: context, size: 12)
        label("Left / right = rider's anatomy", at: CGPoint(x: x, y: 70), context: context, size: 11)
        for (index, name) in jointNames.enumerated() {
            let state = points[name] == nil ? " — hidden / absent" : ""
            label("\(index + 1)  \(name.replacingOccurrences(of: "_", with: " "))\(state)",
                  at: CGPoint(x: x, y: Double(98 + index * 20)), context: context, size: 11)
        }
        for (index, text) in ["Cyan: raw detected joints / bones", "Yellow: selected rider crop",
                              "Pink: provisional reference circles", "2D confidence < 0.2: hidden"].enumerated() {
            label(text, at: CGPoint(x: x, y: Double(495 + index * 17)), context: context, size: 11)
        }
        guard let output=context.makeImage(),let dest=CGImageDestinationCreateWithURL(url as CFURL,UTType.png.identifier as CFString,1,nil) else {
            throw PoseError.invalidInput("Cannot create output PNG.")
        }
        CGImageDestinationAddImage(dest,output,nil)
        guard CGImageDestinationFinalize(dest) else { throw PoseError.invalidInput("PNG export failed.") }
    }
    private static func label(_ text: String, at point: CGPoint, context: CGContext, size: CGFloat = 7) {
        context.saveGState(); context.translateBy(x: point.x,y: point.y); context.scaleBy(x: 1,y: -1)
        let attrs: [NSAttributedString.Key: Any] = [NSAttributedString.Key(kCTFontAttributeName as String): CTFontCreateWithName("Menlo" as CFString,size,nil),
            NSAttributedString.Key(kCTForegroundColorAttributeName as String): CGColor(gray: 1,alpha: 1)]
        let line=CTLineCreateWithAttributedString(NSAttributedString(string:text,attributes:attrs))
        let bounds=CTLineGetBoundsWithOptions(line,[])
        context.setFillColor(CGColor(gray:0,alpha:0.75)); context.fill(bounds.insetBy(dx:-1,dy:-1)); context.textPosition = .zero
        CTLineDraw(line,context); context.restoreGState()
    }
}
