"""
Cross-GPU Streaming System for Direct GPU-to-GPU Data Transfer
Uses electron-speed operations with no pre-allocation
"""
from typing import Dict, List, Optional, Any, Tuple
import threading
import time
import logging
import os
import numpy as np
from pathlib import Path
from electron_speed import max_switch_freq, GATE_DELAY

class CrossGPUStream:
    """
    Stream-based GPU-to-GPU data transfer system optimized for mining
    Uses electron-speed operations and direct streaming
    """
    def __init__(self, gpu_id: int, block_size: int = 256):
        self.gpu_id = gpu_id
        self.block_size = block_size  # SHA-256 block size
        self.switch_freq = max_switch_freq
        self.gate_delay = GATE_DELAY
        
        # Initialize stream directory
        self.stream_path = Path(f"gpu_storage/gpu_{gpu_id}/streams")
        os.makedirs(self.stream_path, exist_ok=True)
        
        # Thread safety
        self.stream_lock = threading.Lock()
        self.active_streams = {}
        
    def write_tensor(self, tensor_array: np.ndarray, data: bytes, target_core: Any) -> str:
        """Stream data to a tensor core with zero-copy operations"""
        stream_id = f"stream_{self.gpu_id}_{int(time.time_ns())}"
        
        with self.stream_lock:
            # Stream directly at electron speed
            transfer_time = len(data) / (self.switch_freq * self.block_size)
            
            # Create stream entry
            self.active_streams[stream_id] = {
                'start_time': time.time_ns(),
                'transfer_time': transfer_time,
                'target_core': target_core.core_id,
                'size': len(data)
            }
            
            # Stream to tensor array without allocation
            tensor_array.reshape(-1)[:len(data)] = np.frombuffer(data, dtype=np.uint8)
            
        return stream_id
        
    def read_tensor(self, stream_id: str) -> Optional[np.ndarray]:
        """Read from a tensor stream"""
        with self.stream_lock:
            if stream_id not in self.active_streams:
                return None
                
            stream_info = self.active_streams[stream_id]
            transfer_complete = (time.time_ns() - stream_info['start_time']) / 1e9 > stream_info['transfer_time']
            
            if not transfer_complete:
                return None
                
            # Stream completed, return and cleanup
            del self.active_streams[stream_id]
            return np.zeros(stream_info['size'], dtype=np.uint8)  # Return empty array, data already in tensor
            
    def get_stream_status(self, stream_id: str) -> Dict[str, Any]:
        """Get status of a stream transfer"""
        with self.stream_lock:
            if stream_id not in self.active_streams:
                return {'status': 'not_found'}
                
            stream_info = self.active_streams[stream_id]
            elapsed = (time.time_ns() - stream_info['start_time']) / 1e9
            progress = min(1.0, elapsed / stream_info['transfer_time'])
            
            return {
                'status': 'completed' if progress >= 1.0 else 'in_progress',
                'progress': progress,
                'target_core': stream_info['target_core'],
                'transfer_time': stream_info['transfer_time'],
                'electron_speed': self.switch_freq
            }
           
                
        # Initialize with empty streams data
        return {
            'streams': {},
            'schema': {
                'streams_table': """
                    CREATE TABLE IF NOT EXISTS streams (
                        stream_id VARCHAR PRIMARY KEY,
                        source_gpu VARCHAR,
                        target_gpu VARCHAR,
                        state VARCHAR,
                        created_at TIMESTAMP,
                        last_active TIMESTAMP,
                        transfer_count INTEGER,
                        total_bytes_transferred BIGINT
                    )""",
                'transfers_table': """
                    CREATE TABLE IF NOT EXISTS transfers (
                        transfer_id VARCHAR PRIMARY KEY,
                        stream_id VARCHAR,
                        transfer_size BIGINT,
                        started_at TIMESTAMP,
                        completed_at TIMESTAMP,
                        state VARCHAR
                    )"""
            }
        }
    def create_stream(self, stream_id: str, source_gpu: str, target_gpu: str) -> Dict[str, Any]:
        """Create a new cross-GPU stream"""
        with self.stream_lock:
            # Check if stream exists
            if stream_id in self.data['streams']:
                return self.data['streams'][stream_id]
            
            # Create new stream
            now = time.time()
            stream_data = {
                'id': stream_id,
                'source_gpu': source_gpu,
                'target_gpu': target_gpu,
                'state': 'initialized',
                'created_at': now,
                'last_active': now,
                'transfer_count': 0,
                'total_bytes_transferred': 0,
                'transfers': {}
            }
            
            self.data['streams'][stream_id] = stream_data
            self._save_storage()
            
            logging.info(f"Created cross-GPU stream {stream_id} from {source_gpu} to {target_gpu}")
            
            return stream_data
            
    def start_transfer(self, stream_id: str, transfer_size: int) -> str:
        """Start a new data transfer on a stream"""
        with self.stream_lock:
            # Check if stream exists
            if stream_id not in self.data['streams']:
                raise ValueError(f"Stream {stream_id} does not exist")
            
            stream = self.data['streams'][stream_id]
            transfer_count = stream['transfer_count']
            transfer_id = f"transfer_{stream_id}_{transfer_count}"
            now = time.time()
            
            # Create transfer record
            with self.transfer_lock:
                transfer_data = {
                    'id': transfer_id,
                    'stream_id': stream_id,
                    'transfer_size': transfer_size,
                    'started_at': now,
                    'state': 'in_progress'
                }
                stream['transfers'][transfer_id] = transfer_data
            
            # Update stream
            stream['transfer_count'] += 1
            stream['last_active'] = now
            self._save_storage()
            
            logging.info(f"Started transfer {transfer_id} on stream {stream_id}")
            return transfer_id
            
    def complete_transfer(self, transfer_id: str) -> None:
        """Mark a transfer as complete and update statistics"""
        with self.transfer_lock:
            # Get transfer info
            transfer = self.con.execute("""
                SELECT stream_id, transfer_size FROM transfers 
                WHERE transfer_id = ? AND state = 'in_progress'
            """, [transfer_id]).fetchone()
            
            if not transfer:
                raise ValueError(f"Transfer {transfer_id} not found or not in progress")
            
            stream_id, transfer_size = transfer
            now = time.time()
            
            # Update transfer status
            self.con.execute("""
                UPDATE transfers 
                SET state = 'completed',
                    completed_at = ?
                WHERE transfer_id = ?
            """, [now, transfer_id])
            
            # Update stream statistics
            with self.stream_lock:
                self.con.execute("""
                    UPDATE streams 
                    SET total_bytes_transferred = total_bytes_transferred + ?,
                        last_active = ?
                    WHERE stream_id = ?
                """, [transfer_size, now, stream_id])
            
            logging.info(f"Completed transfer {transfer_id}")
            
    def get_stream_stats(self, stream_id: str) -> Dict[str, Any]:
        """Get statistics for a specific stream"""
        with self.stream_lock:
            if stream_id not in self.data['streams']:
                raise ValueError(f"Stream {stream_id} does not exist")
            
            stream = self.data['streams'][stream_id]
            now = time.time()
            
            return {
                'transfer_count': stream['transfer_count'],
                'total_bytes_transferred': stream['total_bytes_transferred'],
                'uptime': now - stream['created_at'],
                'last_active_ago': now - stream['last_active']
            }
            
    def cleanup_inactive_streams(self, timeout: float = 300.0) -> List[str]:
        """Remove streams that have been inactive for the specified timeout"""
        current_time = time.time()
        cutoff_time = current_time - timeout
        cleaned_streams = []
        
        with self.stream_lock:
            # Find and remove inactive streams
            stream_ids = list(self.data['streams'].keys())
            for stream_id in stream_ids:
                if self.data['streams'][stream_id]['last_active'] < cutoff_time:
                    del self.data['streams'][stream_id]
                    cleaned_streams.append(stream_id)
                    logging.info(f"Cleaned up inactive stream {stream_id}")
            
            if cleaned_streams:
                self._save_storage()
            
            return cleaned_streams
