from typing import Dict, List, Tuple, Any, Optional
import numpy as np
import threading
import time
import json
from dataclasses import dataclass
from enum import Enum
from .memory import RegisterFile
from .hal.storage_base import VirtualGPUStorageBase, DeviceType



@dataclass
class WarpBarrier:
    """Represents a synchronization barrier for warps"""
    barrier_id: str
    num_warps: int
    arrived: int = 0
    completed: bool = False
    lock: threading.Lock = threading.Lock()
    condition: threading.Condition = threading.Condition()

class ShuffleMode(Enum):
    """Different modes for warp shuffle operations"""
    UP = "up"  # Shuffle up relative to caller
    DOWN = "down"  # Shuffle down relative to caller
    XOR = "xor"  # Butterfly shuffle pattern
    IDX = "idx"  # Direct index-based shuffle
    BCAST = "bcast"  # Broadcast from source lane

class VotingMode(Enum):
    """Different modes for warp voting operations"""
    ALL = "all"  # True if predicate is true for all active threads
    ANY = "any"  # True if predicate is true for any active thread
    BALLOT = "ballot"  # Returns bitmask of true predicates
    COUNT = "count"  # Returns count of true predicates

class WarpState(Enum):
    """Possible states for a warp"""
    READY = "ready"  # Ready to execute
    RUNNING = "running"  # Currently executing
    BLOCKED = "blocked"  # Waiting for synchronization
    YIELDED = "yielded"  # Voluntarily yielded execution
    COMPLETED = "completed"  # Finished execution

class Warp(VirtualGPUStorageBase):
    """Represents a group of threads that execute together with advanced synchronization"""
    
    def __init__(self, warp_id: int, num_threads: int = 32):
        # Initialize local storage
        super().__init__(f"warp_{warp_id}", DeviceType.WARP)
        
        self.warp_id = warp_id
        self.num_threads = min(num_threads, 32)  # Max 32 threads per warp
        self.active_mask = (1 << self.num_threads) - 1  # All threads active initially
        self.predicate_mask = (1 << self.num_threads) - 1  # For predicated execution
        self.registers = [RegisterFile() for _ in range(self.num_threads)]
        self.state = WarpState.READY
        
        # Initialize database schema
        self._init_database()
        
        # Register warp in database
        self._register_warp()
        
    def _init_database(self):
        """Initialize database tables"""
        super()._init_database()  # Initialize base tables first

    def _init_database(self):
        """Initialize database tables"""
        super()._init_database()  # Initialize base tables first
        
        # Warp state table
        self.conn.execute("""
            CREATE TABLE IF NOT EXISTS warps (
                warp_id VARCHAR PRIMARY KEY,
                num_threads INTEGER,
                active_mask BIGINT,
                predicate_mask BIGINT,
                state VARCHAR,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                updated_at TIMESTAMP,
                state_json JSON
            )
        """)
        
        # Barrier table
        self.conn.execute("""
            CREATE TABLE IF NOT EXISTS warp_barriers (
                barrier_id VARCHAR PRIMARY KEY,
                num_warps INTEGER,
                arrived_count INTEGER DEFAULT 0,
                completed BOOLEAN DEFAULT false,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                completed_at TIMESTAMP,
                state_json JSON
            )
        """)
        
        # Register table
        self.conn.execute("""
            CREATE TABLE IF NOT EXISTS warp_registers (
                warp_id VARCHAR,
                thread_id INTEGER,
                register_id INTEGER,
                value BLOB,
                updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                PRIMARY KEY (warp_id, thread_id, register_id)
            )
        """)
        
    def _register_warp(self):
        """Register warp in database"""
        self.conn.execute("""
            INSERT INTO warps (
                warp_id, num_threads, active_mask, predicate_mask,
                state, state_json
            ) VALUES (?, ?, ?, ?, ?, ?)
        """, [
            str(self.warp_id),
            self.num_threads,
            self.active_mask,
            self.predicate_mask,
            self.state.value,
            {"status": "initialized"}
        ])
        self.pc = 0  # Program counter
        
        # Synchronization
        self.barriers: Dict[str, WarpBarrier] = {}
        self.lock = threading.Lock()
        
        # Performance tracking
        self.cycles_executed = 0
        self.last_active_time = time.time()
        
    def get_active_threads(self) -> List[int]:
        """Get indices of currently active threads"""
        return [i for i in range(self.num_threads) 
                if self.active_mask & (1 << i)]
                
    def get_predicated_threads(self) -> List[int]:
        """Get indices of threads that pass predication"""
        return [i for i in range(self.num_threads)
                if (self.active_mask & (1 << i)) and (self.predicate_mask & (1 << i))]
                
    def set_active_mask(self, mask: int):
        """Set which threads are active"""
        with self.lock:
            self.active_mask = mask & ((1 << self.num_threads) - 1)
            
    def set_predicate_mask(self, mask: int):
        """Set predication mask for conditional execution"""
        with self.lock:
            self.predicate_mask = mask & ((1 << self.num_threads) - 1)
            
    def sync(self, barrier_id: str = None):
        """Synchronize all threads in the warp at a barrier"""
        if not barrier_id:
            barrier_id = f"warp_{self.warp_id}_barrier_{time.time_ns()}"
            
        with self.lock:
            if barrier_id not in self.barriers:
                self.barriers[barrier_id] = WarpBarrier(
                    barrier_id=barrier_id,
                    num_warps=1
                )
            
            barrier = self.barriers[barrier_id]
            with barrier.lock:
                barrier.arrived += 1
                if barrier.arrived == barrier.num_warps:
                    barrier.completed = True
                    with barrier.condition:
                        barrier.condition.notify_all()
                else:
                    while not barrier.completed:
                        with barrier.condition:
                            barrier.condition.wait()
                            
    def vote(self, predicate: List[bool], mode: VotingMode = VotingMode.ALL) -> Any:
        """Perform voting operation across threads"""
        active_threads = self.get_predicated_threads()
        if not active_threads:
            return False if mode != VotingMode.BALLOT else 0
            
        if mode == VotingMode.ALL:
            return all(predicate[i] for i in active_threads)
        elif mode == VotingMode.ANY:
            return any(predicate[i] for i in active_threads)
        elif mode == VotingMode.BALLOT:
            return sum(1 << i for i in active_threads if predicate[i])
        elif mode == VotingMode.COUNT:
            return sum(1 for i in active_threads if predicate[i])
            
    def shuffle(self, var: List[Any], mode: ShuffleMode, offset: int) -> List[Any]:
        """Exchange variables between threads using different shuffle patterns"""
        active_threads = self.get_predicated_threads()
        result = list(var)  # Create copy to store results
        
        if mode == ShuffleMode.UP:
            # Shift values up by offset
            for i in active_threads:
                src_idx = (i - offset) % self.num_threads
                if src_idx in active_threads:
                    result[i] = var[src_idx]
                    
        elif mode == ShuffleMode.DOWN:
            # Shift values down by offset
            for i in active_threads:
                src_idx = (i + offset) % self.num_threads
                if src_idx in active_threads:
                    result[i] = var[src_idx]
                    
        elif mode == ShuffleMode.XOR:
            # Butterfly pattern exchange
            for i in active_threads:
                src_idx = i ^ offset  # XOR with offset
                if src_idx < self.num_threads and src_idx in active_threads:
                    result[i] = var[src_idx]
                    
        elif mode == ShuffleMode.IDX:
            # Direct index-based shuffle
            for i in active_threads:
                if offset < self.num_threads and offset in active_threads:
                    result[i] = var[offset]
                    
        elif mode == ShuffleMode.BCAST:
            # Broadcast from source lane
            if offset < self.num_threads and offset in active_threads:
                src_val = var[offset]
                for i in active_threads:
                    result[i] = src_val
                    
        return result
        
    def execute(self, func: callable, *args, **kwargs):
        """Execute a function across all active threads"""
        active_threads = self.get_active_threads()
        results = []
        
        for thread_idx in active_threads:
            # Set up thread context
            ctx = {
                'thread_idx': thread_idx,
                'warp_id': self.warp_id,
                'registers': self.registers[thread_idx]
            }
            
            # Execute thread
            result = func(ctx, *args, **kwargs)
            results.append(result)
            
        return results

class WarpScheduler:
    """Advanced warp scheduler with priority and dependency handling"""
    
    def __init__(self, max_warps: int = 32, max_active_warps: int = 16):
        self.max_warps = max_warps
        self.max_active_warps = max_active_warps
        self.warps: List[Warp] = []
        self.active_warps: Dict[int, bool] = {}
        self.warp_priorities: Dict[int, int] = {}
        self.warp_dependencies: Dict[int, List[int]] = {}
        self.lock = threading.Lock()
        
    def create_warp(self, num_threads: int = 32, priority: int = 0) -> Warp:
        """Create a new warp with specified priority"""
        with self.lock:
            if len(self.warps) >= self.max_warps:
                raise RuntimeError("Maximum number of warps reached")
                
            warp_id = len(self.warps)
            warp = Warp(warp_id, num_threads)
            self.warps.append(warp)
            self.active_warps[warp_id] = True
            self.warp_priorities[warp_id] = priority
            self.warp_dependencies[warp_id] = []
            return warp
            
    def set_warp_priority(self, warp_id: int, priority: int):
        """Set execution priority for a warp"""
        with self.lock:
            if 0 <= warp_id < len(self.warps):
                self.warp_priorities[warp_id] = priority
                
    def add_warp_dependency(self, warp_id: int, depends_on: int):
        """Add execution dependency between warps"""
        with self.lock:
            if 0 <= warp_id < len(self.warps) and 0 <= depends_on < len(self.warps):
                self.warp_dependencies[warp_id].append(depends_on)
                
    def remove_warp_dependency(self, warp_id: int, depends_on: int):
        """Remove execution dependency between warps"""
        with self.lock:
            if 0 <= warp_id < len(self.warps):
                try:
                    self.warp_dependencies[warp_id].remove(depends_on)
                except ValueError:
                    pass
                    
    def suspend_warp(self, warp_id: int):
        """Suspend a warp from execution"""
        with self.lock:
            if 0 <= warp_id < len(self.warps):
                self.active_warps[warp_id] = False
                self.warps[warp_id].state = WarpState.BLOCKED
                
    def resume_warp(self, warp_id: int):
        """Resume a suspended warp"""
        with self.lock:
            if 0 <= warp_id < len(self.warps):
                self.active_warps[warp_id] = True
                self.warps[warp_id].state = WarpState.READY
                
    def synchronize_warps(self, warp_ids: List[int], barrier_id: str = None):
        """Synchronize a group of warps"""
        if not barrier_id:
            barrier_id = f"barrier_{time.time_ns()}"
            
        # Create barrier
        barrier = WarpBarrier(barrier_id=barrier_id, num_warps=len(warp_ids))
        
        # Register barrier with each warp
        for warp_id in warp_ids:
            if 0 <= warp_id < len(self.warps):
                warp = self.warps[warp_id]
                warp.barriers[barrier_id] = barrier
                
        # Wait for all warps to reach barrier
        for warp_id in warp_ids:
            if 0 <= warp_id < len(self.warps):
                self.warps[warp_id].sync(barrier_id)
                
    def schedule(self) -> List[Warp]:
        """Schedule warps for execution based on priority and dependencies"""
        with self.lock:
            ready_warps = []
            
            # Get warps that are ready to execute
            for warp_id, warp in enumerate(self.warps):
                if not self.active_warps.get(warp_id, False):
                    continue
                    
                # Check dependencies
                dependencies_met = all(
                    self.warps[dep_id].state == WarpState.COMPLETED
                    for dep_id in self.warp_dependencies.get(warp_id, [])
                )
                
                if dependencies_met and warp.state == WarpState.READY:
                    ready_warps.append((warp_id, self.warp_priorities.get(warp_id, 0)))
                    
            # Sort by priority (higher numbers = higher priority)
            ready_warps.sort(key=lambda x: x[1], reverse=True)
            
            # Return warps up to max_active_warps
            return [self.warps[warp_id] 
                   for warp_id, _ in ready_warps[:self.max_active_warps]]
                   
    def execute_warps(self, func: callable, *args, **kwargs):
        """Execute function across all active warps with scheduling"""
        results = []
        scheduled_warps = self.schedule()
        
        for warp in scheduled_warps:
            warp.state = WarpState.RUNNING
            result = warp.execute(func, *args, **kwargs)
            results.extend(result)
            warp.last_active_time = time.time()
            warp.cycles_executed += 1
            warp.state = WarpState.READY
            
        return results
