import asyncio
import json
import os
from pathlib import Path
import uuid
import time
from typing import Dict, Any, Optional
import numpy as np
from fastapi import FastAPI, WebSocket
from fastapi.responses import HTMLResponse
from datetime import datetime

# Create FastAPI instance
app = FastAPI()

class VirtualGPUServer:
    def __init__(self):
        self.base_path = Path(__file__).parent / "storage"
        self.vram_path = self.base_path / "vram_blocks"
        self.state_path = self.base_path / "gpu_state"
        self.cache_path = self.base_path / "cache"
        
        # Ensure all storage directories exist
        self.vram_path.mkdir(parents=True, exist_ok=True)
        self.state_path.mkdir(parents=True, exist_ok=True)
        self.cache_path.mkdir(parents=True, exist_ok=True)
        
        # In-memory caches for faster access
        self.vram_cache: Dict[str, Any] = {}
        self.state_cache: Dict[str, Any] = {}
        self.memory_cache: Dict[str, Any] = {}
        
        # Active connections and sessions
        self.active_connections: Dict[str, WebSocket] = {}
        self.active_sessions: Dict[str, Dict[str, Any]] = {}
        self.heartbeat_interval = 5  # seconds
        self.connection_timeout = 30  # seconds
        
        # Performance monitoring
        self.ops_counter = 0
        self.start_time = time.time()

    def _make_json_serializable(self, obj):
        """Convert non-JSON-serializable objects to serializable format"""
        if isinstance(obj, dict):
            return {k: self._make_json_serializable(v) for k, v in obj.items()}
        elif isinstance(obj, list):
            return [self._make_json_serializable(i) for i in obj]
        elif isinstance(obj, tuple):
            return list(obj)
        elif isinstance(obj, (np.ndarray, np.generic)):
            return obj.tolist()
        elif isinstance(obj, (Path, uuid.UUID)):
            return str(obj)
        elif hasattr(obj, '__dict__'):
            # Handle custom objects by converting their __dict__ to serializable format
            return self._make_json_serializable(obj.__dict__)
        elif isinstance(obj, (int, float, str, bool, type(None))):
            return obj
        else:
            # Convert any other types to string representation
            return str(obj)

    async def monitor_connection(self, websocket: WebSocket, session_id: str):
        """Monitor connection health and handle reconnection"""
        try:
            while session_id in self.active_connections:
                try:
                    await asyncio.wait_for(websocket.receive_text(), timeout=self.heartbeat_interval)
                except asyncio.TimeoutError:
                    try:
                        await websocket.send_json({"type": "ping"})
                    except:
                        print(f"Connection lost for session {session_id}")
                        break
                except Exception:
                    break
                await asyncio.sleep(self.heartbeat_interval)
        finally:
            await self.handle_disconnect(session_id)

    async def handle_disconnect(self, session_id: str):
        """Clean up resources when a client disconnects"""
        if session_id in self.active_connections:
            try:
                await self.active_connections[session_id].close()
            except:
                pass
            del self.active_connections[session_id]
        if session_id in self.active_sessions:
            # Save any pending state before removing session
            session_data = self.active_sessions[session_id]
            if session_data.get('pending_state'):
                await self.handle_state_operation({
                    'type': 'save',
                    'component': 'session',
                    'state_id': session_id,
                    'data': session_data['pending_state']
                })
            del self.active_sessions[session_id]

    async def handle_vram_operation(self, operation: dict) -> dict:
        """Handle VRAM read/write operations"""
        try:
            op_type = operation.get('type')
            if not op_type:
                raise ValueError("Missing operation type")
            
            block_id = operation.get('block_id')
            if not block_id:
                raise ValueError("Missing block_id")
            
            data = operation.get('data')
            if data and isinstance(data, (dict, list)):
                data = self._make_json_serializable(data)

            if op_type == 'write':
                if data is None:
                    raise ValueError("Missing data for write operation")
                file_path = self.vram_path / f"{block_id}.npy"
                np.save(file_path, np.array(data))
                self.vram_cache[block_id] = np.array(data)
                return {'status': 'success', 'message': f'Block {block_id} written'}
            
            if op_type == 'read':
                if block_id in self.vram_cache:
                    return {
                        'status': 'success',
                        'data': self.vram_cache[block_id] if isinstance(self.vram_cache[block_id], list) else self.vram_cache[block_id].tolist(),
                        'source': 'cache'
                    }
                
                file_path = self.vram_path / f"{block_id}.npy"
                if file_path.exists():
                    data = np.load(file_path)
                    self.vram_cache[block_id] = np.array(data)
                    return {
                        'status': 'success',
                        'data': data.tolist(),
                        'source': 'disk'
                    }
                return {'status': 'error', 'message': 'Block not found'}
            
            return {'status': 'error', 'message': f'Unknown operation type: {op_type}'}
            
        except ValueError as e:
            return {'status': 'error', 'message': str(e)}
        except Exception as e:
            return {'status': 'error', 'message': f'Operation failed: {str(e)}'}
                
    async def handle_state_operation(self, operation: dict) -> dict:
        """Handle GPU state operations"""
        op_type = operation.get('type')
        component = operation.get('component')
        state_id = operation.get('state_id')
        state_data = operation.get('data')

        file_path = self.state_path / component / f"{state_id}.json"
        
        if op_type == 'save':
            file_path.parent.mkdir(exist_ok=True)
            with open(file_path, 'w') as f:
                json.dump(state_data, f)
            self.state_cache[f"{component}:{state_id}"] = state_data
            return {'status': 'success', 'message': f'State {state_id} saved'}
            
        elif op_type == 'load':
            cache_key = f"{component}:{state_id}"
            if cache_key in self.state_cache:
                return {
                    'status': 'success',
                    'data': self.state_cache[cache_key],
                    'source': 'cache'
                }
                
            if file_path.exists():
                with open(file_path) as f:
                    state_data = json.load(f)
                self.state_cache[cache_key] = state_data
                return {
                    'status': 'success',
                    'data': state_data,
                    'source': 'disk'
                }
                
            return {'status': 'error', 'message': 'State not found'}

    async def handle_cache_operation(self, operation: dict) -> dict:
        """Handle cache operations"""
        op_type = operation.get('type')
        key = operation.get('key')
        data = operation.get('data')
        
        if op_type == 'set':
            self.memory_cache[key] = data
            # Also persist to disk for recovery
            file_path = self.cache_path / f"{key}.json"
            with open(file_path, 'w') as f:
                json.dump(data, f)
            return {'status': 'success', 'message': f'Cache key {key} set'}
            
        elif op_type == 'get':
            if key in self.memory_cache:
                return {
                    'status': 'success',
                    'data': self.memory_cache[key],
                    'source': 'memory'
                }
                
            file_path = self.cache_path / f"{key}.json"
            if file_path.exists():
                with open(file_path) as f:
                    data = json.load(f)
                self.memory_cache[key] = data
                return {
                    'status': 'success',
                    'data': data,
                    'source': 'disk'
                }
                
            return {'status': 'error', 'message': 'Cache key not found'}


    async def handle_connection(self, websocket: WebSocket, connection_type: str):
        """Handle incoming WebSocket connections"""
        # Generate unique session ID
        session_id = str(uuid.uuid4())
        self.active_connections[session_id] = websocket
        self.active_sessions[session_id] = {
            'start_time': time.time(),
            'ops_count': 0,
            'connection_type': connection_type
        }
        
        try:
            while True:
                message = await websocket.receive_json()
                
                # Route operation to appropriate handler
                operation_type = message.get('operation')
                if operation_type == 'vram':
                    response = await self.handle_vram_operation(message)
                elif operation_type == 'state':
                    response = await self.handle_state_operation(message)
                elif operation_type == 'cache':
                    response = await self.handle_cache_operation(message)
                else:
                    response = {
                        'status': 'error',
                        'message': 'Unknown operation type'
                    }
                
                # Update statistics
                self.ops_counter += 1
                self.active_sessions[session_id]['ops_count'] += 1
                
                # Send response
                await websocket.send_json(response)
                
        except Exception as e:
            print(f"WebSocket error: {e}")
        finally:
            # Cleanup on disconnect
            if session_id in self.active_connections:
                del self.active_connections[session_id]
            if session_id in self.active_sessions:
                del self.active_sessions[session_id]

    def get_stats(self) -> dict:
        current_time = time.time()
        uptime = current_time - self.start_time
        ops_per_second = self.ops_counter / uptime if uptime > 0 else 0
        
        return {
            'uptime': uptime,
            'total_operations': self.ops_counter,
            'ops_per_second': ops_per_second,
            'active_connections': len(self.active_connections),
            'vram_cache_size': len(self.vram_cache),
            'state_cache_size': len(self.state_cache),
            'memory_cache_size': len(self.memory_cache)
        }

server = VirtualGPUServer()

@app.get("/", response_class=HTMLResponse)
async def handle_index():
    """Handle HTTP index request"""
    stats = server.get_stats()
    html = f"""
    <!DOCTYPE html>
    <html>
    <head>
        <title>Virtual GPU Server</title>
        <style>
            body {{ font-family: Arial, sans-serif; margin: 40px; }}
            table {{ border-collapse: collapse; width: 100%; margin-top: 20px; }}
            th, td {{ padding: 12px; text-align: left; border-bottom: 1px solid #ddd; }}
            th {{ background-color: #f2f2f2; }}
            .stats {{ background-color: #f9f9f9; padding: 20px; border-radius: 5px; }}
        </style>
    </head>
    <body>
        <h1>Virtual GPU Server Status</h1>
        <div class="stats">
            <h2>Server Statistics</h2>
            <ul>
                <li>Uptime: {stats['uptime']:.2f} seconds</li>
                <li>Total Operations: {stats['total_operations']}</li>
                <li>Operations per Second: {stats['ops_per_second']:.2f}</li>
                <li>Active Connections: {stats['active_connections']}</li>
                <li>VRAM Cache Size: {stats['vram_cache_size']}</li>
                <li>State Cache Size: {stats['state_cache_size']}</li>
                <li>Memory Cache Size: {stats['memory_cache_size']}</li>
            </ul>
        </div>
        <h2>Server Files</h2>
        <iframe src="/files" style="width: 100%; height: 500px; border: none;"></iframe>
    </body>
    </html>
    """
    return HTMLResponse(content=html)

@app.get("/files", response_class=HTMLResponse)
async def handle_files():
    """Handle HTTP files listing request"""
    def format_size(size):
        for unit in ['B', 'KB', 'MB', 'GB']:
            if size < 1024:
                return f"{size:.2f} {unit}"
            size /= 1024
        return f"{size:.2f} TB"

    html = ['<!DOCTYPE html><html><head>',
            '<style>',
            'body { font-family: Arial, sans-serif; margin: 20px; }',
            'table { border-collapse: collapse; width: 100%; }',
            'th, td { padding: 12px; text-align: left; border-bottom: 1px solid #ddd; }',
            'th { background-color: #f2f2f2; }',
            '</style></head><body>',
            '<h2>Server Files</h2>',
            '<table><tr><th>Path</th><th>Size</th><th>Last Modified</th></tr>']

    for root, _, files in os.walk(server.base_path):
        for file in files:
            full_path = Path(root) / file
            rel_path = full_path.relative_to(server.base_path)
            size = format_size(os.path.getsize(full_path))
            mtime = datetime.fromtimestamp(os.path.getmtime(full_path))
            html.append(f'<tr><td>{rel_path}</td><td>{size}</td><td>{mtime}</td></tr>')

    html.extend(['</table></body></html>'])
    return HTMLResponse(content='\n'.join(html))

# WebSocket endpoint
@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept()
    await server.handle_connection(websocket, "gpu")

@app.websocket("/ws/model")
async def websocket_model_endpoint(websocket: WebSocket):
    await websocket.accept()
    await server.handle_connection(websocket, "model")