src/driver.py

"""
GPU Driver Module

This module acts as the interface between a virtual CPU (or external command source)
and the vGPU, handling command queuing and interpretation.
"""

import asyncio
from collections import deque
from enum import Enum
from typing import Dict, Any, Optional, List
from dataclasses import dataclass


class CommandType(Enum):
    """Enumeration of supported GPU commands."""
    CLEAR = "clear"
    DRAW_RECT = "draw_rect"
    DRAW_PIXEL = "draw_pixel"
    DRAW_IMAGE = "draw_image"
    SET_SHADER = "set_shader"
    MATRIX_MULTIPLY = "matrix_multiply"
    VECTOR_OP = "vector_op"
    CREATE_FRAMEBUFFER = "create_framebuffer"
    SET_FRAMEBUFFER = "set_framebuffer"
    LOAD_TEXTURE = "load_texture"


@dataclass
class Command:
    """Represents a single command to be executed by the vGPU."""
    command_id: str
    command_type: CommandType
    parameters: Dict[str, Any]
    priority: int = 0
    timestamp: float = 0.0


class GPUDriver:
    """
    GPU Driver that manages command queues and interfaces with the vGPU.
    
    This class receives commands from external sources (virtual CPU, applications)
    and translates them into tasks that can be processed by the vGPU.
    """
    
    def __init__(self, vgpu=None):
        self.vgpu = vgpu
        
        # Command queue management
        self.command_queue = deque()
        self.command_counter = 0
        
        # Current state
        self.current_framebuffer = None
        self.current_shader = None
        
        # Command processing statistics
        self.commands_processed = 0
        self.commands_failed = 0
        
    def set_vgpu(self, vgpu):
        """Set the vGPU reference."""
        self.vgpu = vgpu
        
    def submit_command(self, command_type: CommandType, parameters: Dict[str, Any], 
                      priority: int = 0) -> str:
        """Submit a command to the GPU driver."""
        command_id = f"cmd_{self.command_counter}"
        self.command_counter += 1
        
        command = Command(
            command_id=command_id,
            command_type=command_type,
            parameters=parameters,
            priority=priority,
            timestamp=asyncio.get_event_loop().time()
        )
        
        # Insert command based on priority (higher priority first)
        if priority > 0:
            # Find insertion point for priority queue
            inserted = False
            for i, existing_cmd in enumerate(self.command_queue):
                if existing_cmd.priority < priority:
                    self.command_queue.insert(i, command)
                    inserted = True
                    break
            if not inserted:
                self.command_queue.append(command)
        else:
            self.command_queue.append(command)
            
        return command_id
        
    async def process_commands(self) -> None:
        """Process all pending commands in the queue."""
        while self.command_queue:
            command = self.command_queue.popleft()
            await self._execute_command(command)
            
    async def _execute_command(self, command: Command) -> None:
        """Execute a single command."""
        try:
            if command.command_type == CommandType.CLEAR:
                await self._handle_clear(command)
            elif command.command_type == CommandType.DRAW_RECT:
                await self._handle_draw_rect(command)
            elif command.command_type == CommandType.DRAW_PIXEL:
                await self._handle_draw_pixel(command)
            elif command.command_type == CommandType.DRAW_IMAGE:
                await self._handle_draw_image(command)
            elif command.command_type == CommandType.SET_SHADER:
                await self._handle_set_shader(command)
            elif command.command_type == CommandType.MATRIX_MULTIPLY:
                await self._handle_matrix_multiply(command)
            elif command.command_type == CommandType.VECTOR_OP:
                await self._handle_vector_op(command)
            elif command.command_type == CommandType.CREATE_FRAMEBUFFER:
                await self._handle_create_framebuffer(command)
            elif command.command_type == CommandType.SET_FRAMEBUFFER:
                await self._handle_set_framebuffer(command)
            elif command.command_type == CommandType.LOAD_TEXTURE:
                await self._handle_load_texture(command)
            else:
                print(f"Unknown command type: {command.command_type}")
                self.commands_failed += 1
                return
                
            self.commands_processed += 1
            
        except Exception as e:
            print(f"Error executing command {command.command_id}: {e}")
            self.commands_failed += 1
            
    async def _handle_clear(self, command: Command) -> None:
        """Handle CLEAR command."""
        if self.vgpu and self.current_framebuffer:
            from vgpu import TaskType
            task_id = self.vgpu.submit_task(
                TaskType.RENDER_CLEAR,
                {
                    "framebuffer_id": self.current_framebuffer,
                    **command.parameters
                }
            )
            
    async def _handle_draw_rect(self, command: Command) -> None:
        """Handle DRAW_RECT command."""
        if self.vgpu and self.current_framebuffer:
            from vgpu import TaskType
            task_id = self.vgpu.submit_task(
                TaskType.RENDER_RECT,
                {
                    "framebuffer_id": self.current_framebuffer,
                    **command.parameters
                }
            )
            
    async def _handle_draw_pixel(self, command: Command) -> None:
        """Handle DRAW_PIXEL command."""
        if self.vgpu and self.current_framebuffer:
            from vgpu import TaskType
            # Convert single pixel to a 1x1 rectangle
            params = command.parameters.copy()
            params.update({
                "framebuffer_id": self.current_framebuffer,
                "width": 1,
                "height": 1
            })
            task_id = self.vgpu.submit_task(TaskType.RENDER_RECT, params)
            
    async def _handle_draw_image(self, command: Command) -> None:
        """Handle DRAW_IMAGE command."""
        if self.vgpu and self.current_framebuffer:
            from vgpu import TaskType
            task_id = self.vgpu.submit_task(
                TaskType.RENDER_IMAGE,
                {
                    "framebuffer_id": self.current_framebuffer,
                    **command.parameters
                }
            )
            
    async def _handle_set_shader(self, command: Command) -> None:
        """Handle SET_SHADER command."""
        shader_id = command.parameters.get("shader_id")
        if shader_id:
            self.current_shader = shader_id
            
    async def _handle_matrix_multiply(self, command: Command) -> None:
        """Handle MATRIX_MULTIPLY command."""
        if self.vgpu:
            from vgpu import TaskType
            task_id = self.vgpu.submit_task(
                TaskType.AI_MATRIX_MULTIPLY,
                command.parameters
            )
            
    async def _handle_vector_op(self, command: Command) -> None:
        """Handle VECTOR_OP command."""
        if self.vgpu:
            from vgpu import TaskType
            task_id = self.vgpu.submit_task(
                TaskType.AI_VECTOR_OP,
                command.parameters
            )
            
    async def _handle_create_framebuffer(self, command: Command) -> None:
        """Handle CREATE_FRAMEBUFFER command."""
        if self.vgpu and self.vgpu.vram:
            width = command.parameters.get("width", 800)
            height = command.parameters.get("height", 600)
            channels = command.parameters.get("channels", 3)
            name = command.parameters.get("name")
            
            framebuffer_id = self.vgpu.vram.create_framebuffer(width, height, channels, name)
            
            # Set as current framebuffer if none is set
            if self.current_framebuffer is None:
                self.current_framebuffer = framebuffer_id
                
    async def _handle_set_framebuffer(self, command: Command) -> None:
        """Handle SET_FRAMEBUFFER command."""
        framebuffer_id = command.parameters.get("framebuffer_id")
        if framebuffer_id and self.vgpu and self.vgpu.vram:
            if self.vgpu.vram.get_framebuffer(framebuffer_id):
                self.current_framebuffer = framebuffer_id
                
    async def _handle_load_texture(self, command: Command) -> None:
        """Handle LOAD_TEXTURE command."""
        if self.vgpu and self.vgpu.vram:
            texture_data = command.parameters.get("texture_data")
            name = command.parameters.get("name")
            
            if texture_data is not None:
                texture_id = self.vgpu.vram.load_texture(texture_data, name)
                
    def get_current_framebuffer(self) -> Optional[str]:
        """Get the current active framebuffer ID."""
        return self.current_framebuffer
        
    def get_current_shader(self) -> Optional[str]:
        """Get the current active shader ID."""
        return self.current_shader
        
    def get_stats(self) -> Dict[str, Any]:
        """Get driver statistics."""
        return {
            "commands_in_queue": len(self.command_queue),
            "commands_processed": self.commands_processed,
            "commands_failed": self.commands_failed,
            "current_framebuffer": self.current_framebuffer,
            "current_shader": self.current_shader
        }
        
    # Convenience methods for common operations
    def clear_screen(self, color: tuple = (0, 0, 0)) -> str:
        """Clear the current framebuffer with the specified color."""
        return self.submit_command(CommandType.CLEAR, {"color": color})
        
    def draw_rectangle(self, x: int, y: int, width: int, height: int, 
                      color: tuple = (255, 255, 255)) -> str:
        """Draw a rectangle on the current framebuffer."""
        return self.submit_command(
            CommandType.DRAW_RECT,
            {"x": x, "y": y, "width": width, "height": height, "color": color}
        )
        
    def draw_pixel(self, x: int, y: int, color: tuple = (255, 255, 255)) -> str:
        """Draw a single pixel on the current framebuffer."""
        return self.submit_command(
            CommandType.DRAW_PIXEL,
            {"x": x, "y": y, "color": color}
        )
        
    def create_framebuffer(self, width: int, height: int, channels: int = 3, 
                          name: Optional[str] = None) -> str:
        """Create a new framebuffer."""
        return self.submit_command(
            CommandType.CREATE_FRAMEBUFFER,
            {"width": width, "height": height, "channels": channels, "name": name}
        )
        
    def set_framebuffer(self, framebuffer_id: str) -> str:
        """Set the active framebuffer."""
        return self.submit_command(
            CommandType.SET_FRAMEBUFFER,
            {"framebuffer_id": framebuffer_id}
        )


if __name__ == "__main__":
    # Test the driver
    async def test_driver():
        driver = GPUDriver()
        
        # Submit some test commands
        driver.create_framebuffer(800, 600)
        driver.clear_screen((255, 0, 0))
        driver.draw_rectangle(100, 100, 200, 150, (0, 255, 0))
        driver.draw_pixel(400, 300, (0, 0, 255))
        
        print(f"Driver stats: {driver.get_stats()}")
        
        # Process commands (without vGPU, they won't actually execute)
        await driver.process_commands()
        
        print(f"Driver stats after processing: {driver.get_stats()}")
    
    asyncio.run(test_driver())