// GPStroke.swift
// Grease Pencil data model — captures 2D pencil input, projects to 3D,
// tessellates into renderable triangle strips.

import Observation
import simd
import Metal

// MARK: - Raw 2D sample from Apple Pencil

struct PencilSample {
    let screenPosition: SIMD2<Float>   // in points
    let pressure: Float                 // 0…1 (force / maxForce)
    let altitude: Float                 // radians, 0 = flat, π/2 = upright
    let azimuth: Float                  // radians, direction pencil points
    let timestamp: Double               // seconds
}

// MARK: - Projected 3D point

struct GPStrokePoint {
    var worldPosition: SIMD3<Float>
    var normal: SIMD3<Float>           // drawing plane normal
    var pressure: Float
    var width: Float                    // world-space stroke width
}

// MARK: - A single stroke (one continuous pencil drag)

@Observable
final class GPStroke: Identifiable {
    let id = UUID()
    var points: [GPStrokePoint] = []
    var color: SIMD4<Float> = SIMD4(0.1, 0.1, 0.1, 1.0)
    var baseWidth: Float = 0.05

    // GPU buffer (rebuilt when points change)
    var vertexBuffer: MTLBuffer?
    var vertexCount: Int = 0

    // MARK: - Add a raw sample and project

    func addSample(
        _ sample: PencilSample,
        projector: StrokeProjector
    ) {
        let worldPos = projector.unproject(screenPoint: sample.screenPosition)
        let width = baseWidth * (0.3 + sample.pressure * 0.7)

        let point = GPStrokePoint(
            worldPosition: worldPos,
            normal: projector.drawingPlaneNormal,
            pressure: sample.pressure,
            width: width
        )
        points.append(point)
    }

    // MARK: - Tessellate to triangle strip

    func tessellate(device: MTLDevice) {
        guard points.count >= 2 else { return }

        var vertices: [MeshVertex] = []

        for i in 0..<points.count {
            let p = points[i]
            let tangent: SIMD3<Float>

            if i == 0 {
                tangent = normalize(points[1].worldPosition - p.worldPosition)
            } else if i == points.count - 1 {
                tangent = normalize(p.worldPosition - points[i-1].worldPosition)
            } else {
                tangent = normalize(points[i+1].worldPosition - points[i-1].worldPosition)
            }

            // Perpendicular in the drawing plane
            let bitangent = normalize(cross(tangent, p.normal)) * p.width * 0.5

            let left  = p.worldPosition - bitangent
            let right = p.worldPosition + bitangent

            // Fade alpha at stroke tips
            let tipFade: Float
            let tipDistance = 3
            if i < tipDistance {
                tipFade = Float(i) / Float(tipDistance)
            } else if i > points.count - 1 - tipDistance {
                tipFade = Float(points.count - 1 - i) / Float(tipDistance)
            } else {
                tipFade = 1.0
            }

            let c = SIMD4<Float>(color.x, color.y, color.z, color.w * tipFade)

            vertices.append(MeshVertex(position: left,  normal: p.normal, color: c))
            vertices.append(MeshVertex(position: right, normal: p.normal, color: c))
        }

        vertexCount = vertices.count
        if !vertices.isEmpty {
            vertexBuffer = device.makeBuffer(
                bytes: vertices,
                length: MemoryLayout<MeshVertex>.stride * vertices.count,
                options: .storageModeShared
            )
        }
    }
}

// MARK: - Stroke Projector

/// Projects 2D screen coordinates into 3D world space on a drawing plane.
struct StrokeProjector {
    let viewMatrix: simd_float4x4
    let projectionMatrix: simd_float4x4
    let viewportSize: SIMD2<Float>
    let drawingPlaneNormal: SIMD3<Float>
    let drawingPlanePoint: SIMD3<Float>    // point on the plane

    /// Unproject a screen point onto the drawing plane.
    func unproject(screenPoint: SIMD2<Float>) -> SIMD3<Float> {
        // Normalize to clip space [-1, 1]
        let nx = (2.0 * screenPoint.x / viewportSize.x) - 1.0
        let ny = 1.0 - (2.0 * screenPoint.y / viewportSize.y)

        let invProj = projectionMatrix.inverse
        let invView = viewMatrix.inverse

        // Near and far points in clip space
        let nearClip = SIMD4<Float>(nx, ny, 0, 1)
        let farClip  = SIMD4<Float>(nx, ny, 1, 1)

        // To eye space
        var nearEye = invProj * nearClip
        nearEye /= nearEye.w
        var farEye = invProj * farClip
        farEye /= farEye.w

        // To world space
        let nearWorld = (invView * nearEye).xyz
        let farWorld  = (invView * farEye).xyz

        // Ray
        let rayOrigin = nearWorld
        let rayDir = normalize(farWorld - nearWorld)

        // Intersect with drawing plane
        let denom = dot(drawingPlaneNormal, rayDir)
        if abs(denom) < 1e-6 {
            // Ray parallel to plane — return plane point as fallback
            return drawingPlanePoint
        }
        let t = dot(drawingPlanePoint - rayOrigin, drawingPlaneNormal) / denom
        return rayOrigin + rayDir * max(t, 0)
    }
}

// MARK: - Layer

@Observable
final class GPLayer: Identifiable {
    let id = UUID()
    var name: String
    var strokes: [GPStroke] = []
    var isVisible: Bool = true
    var opacity: Float = 1.0
    var blendMode: GPBlendMode = .normal

    // Per-frame strokes (for animation)
    var frameStrokes: [Int: [GPStroke]] = [:]

    init(name: String) {
        self.name = name
    }
}

enum GPBlendMode: String, CaseIterable {
    case normal = "Normal"
    case multiply = "Multiply"
    case add = "Add"
}

// MARK: - Grease Pencil Object Data

@Observable
final class GPObjectData {
    var layers: [GPLayer] = [GPLayer(name: "Layer")]
    var activeLayerIndex: Int = 0

    var activeLayer: GPLayer? {
        guard layers.indices.contains(activeLayerIndex) else { return nil }
        return layers[activeLayerIndex]
    }

    var currentFrame: Int = 1

    // Onion skinning
    var onionSkinEnabled: Bool = true
    var onionSkinBefore: Int = 3
    var onionSkinAfter: Int = 1

    func addLayer(name: String) {
        layers.append(GPLayer(name: name))
        activeLayerIndex = layers.count - 1
    }

    func removeLayer(at index: Int) {
        guard layers.count > 1 else { return }
        layers.remove(at: index)
        activeLayerIndex = min(activeLayerIndex, layers.count - 1)
    }
}