import { useEffect, useState, Suspense, useMemo, Component, ReactNode, useRef } from 'react';
import { Canvas, useFrame } from '@react-three/fiber';
import { Text, Instances, Instance, Html, CameraControls } from '@react-three/drei';
import * as THREE from 'three';
import { layoutApi, Floor3DData } from '../lib/layoutApi';
import { Loader2, ArrowLeft, Maximize, MousePointer2 } from 'lucide-react';
import { Link } from 'react-router-dom';

// Colors
const COLORS = {
    VEG: '#4ade80',    // green-400
    FLOWER: '#a855f7', // purple-500
    DRY: '#f59e0b',    // amber-500
    CURE: '#78716c',   // stone-500
    EMPTY_SLOT: '#374151', // gray-700
    ROOM_FLOOR: '#1f2937', // gray-800
    ROOM_WALL: '#374151',  // gray-700
};

// Define State interface
interface ErrorBoundaryState {
    hasError: boolean;
    error: Error | null;
}

class ErrorBoundary extends Component<{ children: ReactNode }, ErrorBoundaryState> {
    constructor(props: { children: ReactNode }) {
        super(props);
        this.state = { hasError: false, error: null };
    }

    static getDerivedStateFromError(error: Error): ErrorBoundaryState {
        return { hasError: true, error };
    }

    render() {
        if (this.state.hasError) {
            console.error("3D Viewer Error Caught:", this.state.error);
            return (
                <Html center>
                    <div className="bg-red-900/90 text-white p-4 rounded border border-red-500 max-w-md">
                        <h3 className="font-bold mb-2">3D Rendering Error</h3>
                        <p className="text-sm font-mono break-words">{this.state.error?.message}</p>
                    </div>
                </Html>
            );
        }
        return this.props.children;
    }
}

function PlantInstances({ positions, onPlantClick }: { positions: any[], onPlantClick: (p: any) => void }) {
    if (!positions || !Array.isArray(positions) || positions.length === 0) return null;

    // Safety filter
    const safePositions = positions.filter(p => p && typeof p.x === 'number' && typeof p.y === 'number' && typeof p.z === 'number');

    const plants = safePositions.filter(p => p.plant);
    const emptySlots = safePositions.filter(p => !p.plant);

    // Limit instances to prevent crash on huge datasets just in case
    if (safePositions.length > 5000) {
        console.warn('Too many positions to render:', safePositions.length);
        return null; // Or return a subset
    }

    return (
        <group>
            {/* Active Plants - Render as Cylinders (Stalks/Canopy) for visibility */}
            {plants.length > 0 && (
                <Instances
                    range={plants.length}
                    onClick={(e) => {
                        e.stopPropagation();
                        const index = e.instanceId;
                        if (index !== undefined && plants[index]) {
                            onPlantClick(plants[index]);
                        }
                    }}
                >
                    {/* Make plants much taller and visible: Cylinder radius 0.2, height 0.8 */}
                    <cylinderGeometry args={[0.2, 0.1, 0.8, 8]} />
                    <meshStandardMaterial />
                    {plants.map((pos, i) => (
                        <Instance
                            key={pos.id || i}
                            // Lift y by 0.4 (half height) so it sits ON the shelf
                            position={[pos.x, pos.y + 0.4, pos.z]}
                            color={pos.plant?.stage === 'FLOWER' ? COLORS.FLOWER :
                                pos.plant?.stage === 'DRYING' ? COLORS.DRY :
                                    pos.plant?.stage === 'CURE' ? COLORS.CURE : COLORS.VEG}
                        />
                    ))}
                </Instances>
            )}

            {/* Empty Slots - Keep subtle but visible */}
            {emptySlots.length > 0 && (
                <Instances range={emptySlots.length}>
                    <cylinderGeometry args={[0.05, 0.05, 0.1, 8]} />
                    <meshStandardMaterial color={COLORS.EMPTY_SLOT} opacity={0.3} transparent />
                    {emptySlots.map((pos, i) => (
                        <Instance
                            key={pos.id || i}
                            position={[pos.x, pos.y, pos.z]}
                        />
                    ))}
                </Instances>
            )}
        </group>
    );
}

function FacilityScene({ data, onSelectPlant, targetView, setControls }: {
    data: Floor3DData,
    onSelectPlant: (p: any) => void,
    targetView: { x: number, y: number, z: number, zoom?: boolean } | null,
    setControls: (c: CameraControls | null) => void
}) {
    const controlsRef = useRef<CameraControls>(null);
    const [keysPressed, setKeysPressed] = useState<Record<string, boolean>>({});

    // Register controls
    useEffect(() => {
        if (controlsRef.current) {
            setControls(controlsRef.current);
        }
    }, [controlsRef, setControls]);

    // Handle View Navigation
    useEffect(() => {
        if (targetView && controlsRef.current) {
            const { x, y, z, zoom } = targetView;
            // Smoothly move camera to look at the target room
            // Position camera slightly offset from the target for a good view
            const dist = zoom ? 15 : 40;
            const height = zoom ? 15 : 40;

            controlsRef.current.setLookAt(
                x + dist, height, z + dist, // Camera Position
                x, 0, z,                    // Target Position
                true                        // Enable transition
            );
        }
    }, [targetView]);

    // Keyboard Controls Logic
    useEffect(() => {
        const handleKeyDown = (e: KeyboardEvent) => setKeysPressed(k => ({ ...k, [e.code]: true }));
        const handleKeyUp = (e: KeyboardEvent) => setKeysPressed(k => ({ ...k, [e.code]: false }));

        window.addEventListener('keydown', handleKeyDown);
        window.addEventListener('keyup', handleKeyUp);
        return () => {
            window.removeEventListener('keydown', handleKeyDown);
            window.removeEventListener('keyup', handleKeyUp);
        };
    }, []);

    useFrame((_, delta) => {
        if (!controlsRef.current) return;

        const speed = 20 * delta; // Movement speed

        // WASD / Arrow Keys Navigation
        if (keysPressed['KeyW'] || keysPressed['ArrowUp']) {
            controlsRef.current.forward(speed, true);
        }
        if (keysPressed['KeyS'] || keysPressed['ArrowDown']) {
            controlsRef.current.forward(-speed, true);
        }
        if (keysPressed['KeyA'] || keysPressed['ArrowLeft']) {
            controlsRef.current.truck(-speed, 0, true);
        }
        if (keysPressed['KeyD'] || keysPressed['ArrowRight']) {
            controlsRef.current.truck(speed, 0, true);
        }
    });

    // Process data into flat list of objects for rendering
    const roomMeshes = useMemo(() => {
        return data.rooms.map(room => {
            // Room creates a floor area
            return (
                <group key={room.id} position={[room.posX, 0, room.posY]}>
                    {/* Floor Label */}
                    <Text
                        position={[room.width / 2, 0.1, room.height / 2]}
                        rotation={[-Math.PI / 2, 0, 0]}
                        fontSize={Math.min(room.width, room.height) / 5} // Scale text to room
                        color="white"
                        fillOpacity={0.8}
                        anchorX="center"
                        anchorY="middle"
                    >
                        {room.name}
                    </Text>

                    {/* Floor Plane - Darker for contrast */}
                    <mesh
                        rotation={[-Math.PI / 2, 0, 0]}
                        position={[room.width / 2, 0, room.height / 2]}
                    >
                        <planeGeometry args={[room.width, room.height]} />
                        <meshStandardMaterial color={COLORS.ROOM_FLOOR} metalness={0.2} roughness={0.8} />
                    </mesh>

                    {/* Room Border/Walls */}
                    <lineSegments position={[room.width / 2, 0.1, room.height / 2]} rotation={[-Math.PI / 2, 0, 0]}>
                        <edgesGeometry args={[new THREE.PlaneGeometry(room.width, room.height)]} />
                        <lineBasicMaterial color="#4b5563" linewidth={2} />
                    </lineSegments>

                    {/* Render sections and positions */}
                    {room.sections.map(section => {
                        // Calculate positions for this section
                        const positions = section.positions.map(pos => {
                            // Simple grid layout logic for position coordinates relative to section
                            // Assuming typical rack dimensions
                            const spacing = 0.5;
                            const x = section.posX + (pos.column * spacing);
                            const z = section.posY + (pos.row * spacing); // Z is depth/row
                            const y = 0.5 + (pos.tier * 0.5); // Y is height/tier

                            return { ...pos, x, y, z };
                        });

                        return (
                            <group key={section.id}>
                                {/* Section Label - Lifted higher */}
                                <Text
                                    position={[
                                        section.posX + (section.width / 2),
                                        3,
                                        section.posY + (section.height / 2)
                                    ]}
                                    fontSize={0.8}
                                    color="#9ca3af"
                                >
                                    {section.code}
                                </Text>
                                <PlantInstances positions={positions} onPlantClick={onSelectPlant} />
                            </group>
                        );
                    })}
                </group>
            );
        });
    }, [data, onSelectPlant]);

    return (
        <>
            <ambientLight intensity={1.5} />
            <pointLight position={[20, 30, 20]} intensity={2} />
            <directionalLight position={[-10, 50, -10]} intensity={1} castShadow />

            <group position={[-data.floor.width / 2, 0, -data.floor.height / 2]}>
                {roomMeshes}
            </group>

            {/* Changed from OrbitControls to CameraControls for better programmatic control */}
            <CameraControls
                ref={controlsRef}
                minDistance={5}
                maxDistance={150}
                dollySpeed={1}
                truckSpeed={2}
                infinityDolly={false}
                makeDefault
            />

            <gridHelper args={[200, 100, 0x444444, 0x111111]} position={[0, -0.1, 0]} />
            <axesHelper args={[5]} />
        </>
    );
}

export default function Facility3DViewerPage() {
    const [status, setStatus] = useState('Initializing...');
    const [floorData, setFloorData] = useState<Floor3DData | null>(null);
    const [selectedPlant, setSelectedPlant] = useState<any | null>(null);
    const [targetView, setTargetView] = useState<{ x: number, y: number, z: number, zoom?: boolean } | null>(null);
    const [controls, setControls] = useState<CameraControls | null>(null);

    useEffect(() => {
        loadData();
    }, []);

    async function loadData() {
        setStatus('Loading layout...');
        try {
            const props = await layoutApi.getProperties();
            if (props[0]?.buildings[0]?.floors[0]) {
                const floorId = props[0].buildings[0].floors[0].id;
                setStatus('Fetching 3D assets...');
                const data = await layoutApi.getFloor3D(floorId);
                setFloorData(data);
                setStatus('');
            } else {
                setStatus('No floor layout found');
            }
        } catch (err) {
            setStatus('Error: ' + (err as Error).message);
        }
    }

    const focusRoom = (room: any) => {
        // Calculate room center relative to the scene group offset
        // Scene group is offset by [-floor.width/2, 0, -floor.height/2]
        if (!floorData) return;

        const offsetX = -floorData.floor.width / 2;
        const offsetZ = -floorData.floor.height / 2;

        const centerX = room.posX + (room.width / 2) + offsetX;
        const centerZ = room.posY + (room.height / 2) + offsetZ;

        setTargetView({ x: centerX, y: 0, z: centerZ, zoom: true });
    };

    const resetView = () => {
        setTargetView({ x: 0, y: 0, z: 0, zoom: false });
    };

    return (
        <div className="h-screen w-full relative bg-gray-950 text-white overflow-hidden">
            {/* Header Overlay */}
            <div className="absolute top-0 left-0 right-0 p-4 z-10 flex items-center justify-between bg-gradient-to-b from-black/80 to-transparent pointer-events-none">
                <div className="pointer-events-auto flex items-center gap-4">
                    <Link to="/metrc" className="btn btn-ghost btn-sm text-white">
                        <ArrowLeft size={16} /> Back
                    </Link>
                    <div>
                        <h1 className="text-xl font-bold flex items-center gap-2">
                            Facility Viewer 3D
                            <span className="badge badge-accent text-xs">BETA</span>
                        </h1>
                        <p className="text-xs text-gray-400">
                            {floorData ? `${floorData.floor.name} • ${floorData.stats.occupiedPositions} Plants` : 'Loading...'}
                        </p>
                    </div>
                </div>

                {/* Legend */}
                <div className="flex gap-4 text-xs bg-black/50 p-2 rounded-lg backdrop-blur pointer-events-auto">
                    <div className="flex items-center gap-1"><div className="w-2 h-2 rounded-full bg-[#4ade80]"></div> Veg</div>
                    <div className="flex items-center gap-1"><div className="w-2 h-2 rounded-full bg-[#a855f7]"></div> Flower</div>
                    <div className="flex items-center gap-1"><div className="w-2 h-2 rounded-full bg-[#f59e0b]"></div> Dry</div>
                    <div className="flex items-center gap-1"><div className="w-2 h-2 rounded-full bg-[#374151]"></div> Empty</div>
                </div>
            </div>

            {/* Quick Nav Sidebar */}
            {floorData && (
                <div className="absolute top-24 left-4 z-10 w-48 space-y-2 pointer-events-auto animate-in slide-in-from-left-4">
                    <div className="bg-black/40 backdrop-blur rounded-lg p-3 border border-gray-800">
                        <div className="flex items-center justify-between mb-2">
                            <span className="text-xs font-bold text-gray-400 uppercase tracking-wider">Navigation</span>
                            <button onClick={resetView} className="hover:text-accent transition-colors">
                                <Maximize size={14} />
                            </button>
                        </div>
                        <div className="space-y-1 max-h-[60vh] overflow-y-auto">
                            {floorData.rooms.map(room => (
                                <button
                                    key={room.id}
                                    onClick={() => focusRoom(room)}
                                    className="w-full text-left text-sm px-2 py-1.5 rounded hover:bg-gray-800 transition-colors flex items-center justify-between group"
                                >
                                    <span className="truncate">{room.name}</span>
                                    <MousePointer2 size={12} className="opacity-0 group-hover:opacity-100 text-accent transition-opacity" />
                                </button>
                            ))}
                        </div>
                    </div>

                    {/* Controls Help */}
                    <div className="bg-black/40 backdrop-blur rounded-lg p-3 border border-gray-800 text-xs text-gray-500">
                        <div className="font-bold mb-1 text-gray-400">Controls</div>
                        <div className="grid grid-cols-2 gap-x-2 gap-y-1">
                            <span>Rotate</span> <span className="text-right text-gray-300">Left Click</span>
                            <span>Pan</span> <span className="text-right text-gray-300">Right Click</span>
                            <span>Zoom</span> <span className="text-right text-gray-300">Scroll</span>
                            <span>Move</span> <span className="text-right text-gray-300">W A S D</span>
                        </div>
                    </div>
                </div>
            )}

            {/* Error/Status Overlay */}
            {status && (
                <div className="absolute inset-0 flex items-center justify-center z-50 bg-black/80 backdrop-blur-sm">
                    <div className="text-center">
                        <Loader2 size={32} className="animate-spin text-accent mx-auto mb-2" />
                        <p className="text-gray-300">{status}</p>
                        {status.includes('Error') && (
                            <button onClick={() => window.location.reload()} className="btn btn-sm btn-outline mt-4">
                                Retry
                            </button>
                        )}
                    </div>
                </div>
            )}

            {/* Selection Overlay */}
            {selectedPlant && (
                <div className="absolute bottom-4 right-4 z-20 w-80 bg-gray-900/90 border border-gray-700 rounded-lg p-4 shadow-xl backdrop-blur-md animate-in slide-in-from-right-10 pointer-events-auto">
                    <div className="flex justify-between items-start mb-2">
                        <h3 className="font-bold text-accent">{selectedPlant.plant.tagNumber}</h3>
                        <button
                            onClick={() => setSelectedPlant(null)}
                            className="text-gray-400 hover:text-white"
                        >
                            ×
                        </button>
                    </div>
                    <div className="space-y-2 text-sm">
                        <div className="flex justify-between">
                            <span className="text-gray-400">Strain:</span>
                            <span>{selectedPlant.plant.strain || 'Unknown'}</span>
                        </div>
                        <div className="flex justify-between">
                            <span className="text-gray-400">Stage:</span>
                            <span className="badge badge-accent bg-opacity-20">{selectedPlant.plant.stage || 'N/A'}</span>
                        </div>
                        <div className="flex justify-between">
                            <span className="text-gray-400">Batch:</span>
                            <span>{selectedPlant.plant.batchName || '-'}</span>
                        </div>
                        <div className="flex justify-between">
                            <span className="text-gray-400">Location:</span>
                            <span className="font-mono text-xs">R{selectedPlant.row} T{selectedPlant.tier} S{selectedPlant.slot}</span>
                        </div>
                    </div>
                </div>
            )}

            <Canvas camera={{ position: [20, 30, 20], fov: 60 }} shadows>
                <ErrorBoundary>
                    <Suspense fallback={<Html center>Loading 3D Scene...</Html>}>
                        {floorData && (
                            <FacilityScene
                                data={floorData}
                                onSelectPlant={setSelectedPlant}
                                targetView={targetView}
                                setControls={setControls}
                            />
                        )}
                    </Suspense>
                </ErrorBoundary>
            </Canvas>
        </div>
    );
}
