import SwiftUI
import RealityKit
import Vision
import ARKit
import CoreML

struct ImmersiveView: View {
    @State private var arkitSession: ARKitSession?
    @State private var detectionResults: [VNRecognizedObjectObservation] = []
    @State private var contentEntity = Entity()
    @State private var boundingBoxEntities: [UUID: ModelEntity] = [:]
    @State private var worldTrackingProvider = WorldTrackingProvider()
    @State private var cameraFrameProvider: CameraFrameProvider?
    @State private var isTrackingReady = false

    var body: some View {
        RealityView { content in
            print("🎥 Initializing RealityKit scene...")
            
            // Add a clear reference point
            let testSphere = ModelEntity(mesh: .generateSphere(radius: 0.5))
            testSphere.position = [0, 0.5, -2.0] // Further in front of camera
            testSphere.model?.materials = [SimpleMaterial(color: .red, isMetallic: false)]
            content.add(testSphere)
            
            content.add(contentEntity)

            Task {
                await setupARKitSession()
            }
        } update: { content in
            Task { @MainActor in
                if isTrackingReady {
                    await updateBoundingBoxes()
                }
            }
        }
    }

    // MARK: - Setup ARKit & Camera Frame Provider
    private func setupARKitSession() async {
        do {
            guard CameraFrameProvider.isSupported else {
                print("🚨 CameraFrameProvider is not supported on this device.")
                return
            }

            let cameraProvider = CameraFrameProvider()
            let session = ARKitSession()
            
            // ✅ Explicitly start world tracking provider
            worldTrackingProvider = WorldTrackingProvider()
            
            try await session.run([cameraProvider, worldTrackingProvider])
            
            self.arkitSession = session
            self.cameraFrameProvider = cameraProvider

            print("✅ ARKit Session started. Waiting for world tracking to initialize...")

            // ✅ Wait until world tracking is fully initialized
            for _ in 0..<5 {
                try await Task.sleep(nanoseconds: 1_000_000_000) // Wait 1 second
                if await worldTrackingProvider.allAnchors != nil {
                    self.isTrackingReady = true
                    print("✅ World Tracking is now ready!")
                    break
                }
                print("⏳ Waiting for world tracking to initialize...")
            }

            guard self.isTrackingReady else {
                print("🚨 World tracking failed to initialize.")
                return
            }

            // ✅ Configure Camera Processing
            let formats = CameraVideoFormat.supportedVideoFormats(for: .main, cameraPositions: [.left])
            let highResolutionFormat = formats.max { $0.frameSize.height < $1.frameSize.height }

            guard let highResolutionFormat,
                  let cameraFrameUpdates = cameraProvider.cameraFrameUpdates(for: highResolutionFormat) else {
                print("🚨 Failed to get camera frame updates.")
                return
            }

            print("🎥 Camera processing started.")

            for await cameraFrame in cameraFrameUpdates {
                if let sample = cameraFrame.sample(for: .left) {
                    print("📷 Captured frame from main camera")
                    processFrame(sample.pixelBuffer)
                }
            }
        } catch {
            print("🚨 Error initializing ARKit session: \(error)")
        }
    }

    // MARK: - Process Frame with YOLOv3
    private func processFrame(_ pixelBuffer: CVPixelBuffer) {
        guard let model = try? VNCoreMLModel(for: YOLOv3().model) else {
            print("🚨 Failed to load YOLOv3 model.")
            return
        }

        let request = VNCoreMLRequest(model: model) { request, _ in
            guard let results = request.results as? [VNRecognizedObjectObservation] else {
                print("🚨 No objects detected.")
                return
            }

            DispatchQueue.main.async {
                self.detectionResults = results
                print("🔍 Detected \(results.count) object(s)")
                for result in results {
                    if let label = result.labels.first {
                        print("🛑 Object: \(label.identifier) | Confidence: \(label.confidence * 100)%")
                    }
                }
            }
        }

        let handler = VNImageRequestHandler(cvPixelBuffer: pixelBuffer)
        do {
            try handler.perform([request])
        } catch {
            print("🚨 Error performing object detection: \(error)")
        }
    }

    // MARK: - Update Bounding Boxes in World Space
    private func updateBoundingBoxes() async {
        guard isTrackingReady else {
            print("🚨 World tracking is not ready yet.")
            return
        }

        guard let cameraTransform = await worldTrackingProvider.allAnchors?.last?.originFromAnchorTransform else {
            print("⏳ Waiting for a valid camera transform...")
            return
        }

        let currentIds = Set(detectionResults.map { $0.uuid })

        print("📡 Updating bounding boxes (\(currentIds.count) detected objects).")

        // ✅ Remove outdated bounding boxes
        boundingBoxEntities.keys.filter { !currentIds.contains($0) }
            .forEach { id in
                print("❌ Removing outdated bounding box: \(id)")
                boundingBoxEntities[id]?.removeFromParent()
                boundingBoxEntities.removeValue(forKey: id)
            }

        // ✅ Add new bounding boxes
        for observation in detectionResults {
            if let entity = await createOrUpdateBoundingBox(for: observation, cameraTransform: cameraTransform) {
                if boundingBoxEntities[observation.uuid] == nil {
                    contentEntity.addChild(entity)
                    boundingBoxEntities[observation.uuid] = entity
                    print("✅ Added bounding box for detected object: \(observation.labels.first?.identifier ?? "Unknown")")
                }
            }
        }
    }

    // MARK: - Convert Bounding Box to World Space
    private func createOrUpdateBoundingBox(for observation: VNRecognizedObjectObservation, cameraTransform: simd_float4x4) async -> ModelEntity? {
        let boundingBox = observation.boundingBox
        let entity = boundingBoxEntities[observation.uuid] ?? ModelEntity()

        // ✅ Estimate depth based on world tracking data
        let estimatedDepth: Float = 1.5
        let cameraPosition = cameraTransform.columns.3

        let worldPosition = SIMD3(
            x: Float(boundingBox.midX - 0.5) * estimatedDepth + cameraPosition.x,
            y: Float(0.5 - boundingBox.midY) * estimatedDepth + cameraPosition.y,
            z: -estimatedDepth + cameraPosition.z
        )

        let boxSize = Float(min(boundingBox.width, boundingBox.height)) * estimatedDepth
        let newMesh = MeshResource.generateBox(width: boxSize, height: boxSize, depth: 0.02)
        let material = UnlitMaterial(color: .green.withAlphaComponent(0.5))
        entity.model = ModelComponent(mesh: newMesh, materials: [material])
        entity.position = worldPosition

        print("📌 Bounding box position: \(entity.position)")

        return entity
    }
}