#!/usr/bin/env python3
# SPDX-License-Identifier: BSD-3-Clause
# Copyright (c) 2026 Open-IOV contributors
"""
guest_test_agent.py - Runs inside guest VM, executes tests, reports results
Communication via vsock, serial, or HTTP
"""

import sys
import os
import json
import time
import subprocess
import threading
import socket
import argparse
from datetime import datetime
from pathlib import Path

class GuestTestAgent:
    def __init__(self, listen_port=9999, report_url=None, vsock_cid=3, vsock_port=9999):
        self.listen_port = listen_port
        self.report_url = report_url
        self.vsock_cid = vsock_cid
        self.vsock_port = vsock_port
        self.running = False
        self.results = []
        
    def start_server(self):
        """Start command listener (TCP or vsock)"""
        self.running = True
        
        # Try vsock first (preferred for VM-host communication)
        try:
            import socket
            if hasattr(socket, 'AF_VSOCK'):
                self.server = socket.socket(socket.AF_VSOCK, socket.SOCK_STREAM)
                self.server.bind((self.vsock_cid, self.vsock_port))
                self.server.listen(5)
                print(f"[{datetime.now()}] Guest agent listening on vsock CID:{self.vsock_cid} PORT:{self.vsock_port}")
            else:
                raise ImportError("vsock not available")
        except Exception:
            # Fallback to TCP
            self.server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            self.server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
            self.server.bind(('0.0.0.0', self.listen_port))
            self.server.listen(5)
            print(f"[{datetime.now()}] Guest agent listening on TCP port {self.listen_port}")
        
        while self.running:
            try:
                conn, addr = self.server.accept()
                threading.Thread(target=self._handle_connection, args=(conn,), daemon=True).start()
            except Exception as e:
                if self.running:
                    print(f"[{datetime.now()}] Server error: {e}")
    
    def _handle_connection(self, conn):
        """Handle incoming test commands"""
        try:
            data = conn.recv(4096).decode('utf-8')
            if not data:
                return
            
            cmd = json.loads(data)
            action = cmd.get('action')
            test_id = cmd.get('test_id', 'unknown')
            
            print(f"[{datetime.now()}] Received command: {action} (test_id: {test_id})")
            
            if action == 'ping':
                result = {'status': 'ok', 'agent': 'guest_test_agent', 'version': '1.0'}
            elif action == 'run_test':
                result = self._run_test(cmd.get('test'), cmd.get('args', {}))
            elif action == 'run_suite':
                result = self._run_suite(cmd.get('suite'))
            elif action == 'get_results':
                result = {'results': self.results}
            elif action == 'shutdown':
                self.running = False
                result = {'status': 'shutting_down'}
            else:
                result = {'error': f'Unknown action: {action}'}
            
            result['test_id'] = test_id
            conn.send(json.dumps(result).encode('utf-8'))
            
        except Exception as e:
            conn.send(json.dumps({'error': str(e)}).encode('utf-8'))
        finally:
            conn.close()
    
    def _run_test(self, test_name, args):
        """Run a single test"""
        test_method = getattr(self, f'test_{test_name}', None)
        if not test_method:
            return {'status': 'error', 'message': f'Unknown test: {test_name}'}
        
        start = time.time()
        try:
            result = test_method(**args)
            result['duration'] = time.time() - start
            result['timestamp'] = datetime.now().isoformat()
            self.results.append(result)
            return result
        except Exception as e:
            return {'status': 'error', 'message': str(e), 'duration': time.time() - start}
    
    def _run_suite(self, suite_name):
        """Run a test suite"""
        suite_method = getattr(self, f'suite_{suite_name}', None)
        if not suite_method:
            return {'status': 'error', 'message': f'Unknown suite: {suite_name}'}
        
        return suite_method()
    
    # ============================================================
    # INDIVIDUAL TESTS
    # ============================================================
    
    def test_nvidia_smi(self):
        """Run nvidia-smi and verify GPU detection"""
        try:
            out = subprocess.check_output(['nvidia-smi'], text=True, timeout=30)
            return {
                'test': 'nvidia_smi',
                'status': 'pass' if 'NVIDIA' in out else 'fail',
                'output': out[:2000]
            }
        except subprocess.CalledProcessError as e:
            return {'test': 'nvidia_smi', 'status': 'fail', 'error': str(e), 'output': e.output}
        except Exception as e:
            return {'test': 'nvidia_smi', 'status': 'error', 'message': str(e)}
    
    def test_cuda_vectoradd(self):
        """Compile and run CUDA vectorAdd sample"""
        try:
            # Find CUDA samples
            cuda_samples = Path('/usr/local/cuda/samples')
            if not cuda_samples.exists():
                cuda_samples = Path('/opt/cuda/samples')
            
            vecadd_dir = cuda_samples / '1_Utilities' / 'vectorAdd'
            if not vecadd_dir.exists():
                return {'test': 'cuda_vectoradd', 'status': 'skip', 'message': 'vectorAdd sample not found'}
            
            # Build
            subprocess.check_call(['make'], cwd=vecadd_dir, timeout=60)
            
            # Run
            out = subprocess.check_output(['./vectorAdd'], cwd=vecadd_dir, text=True, timeout=30)
            
            return {
                'test': 'cuda_vectoradd',
                'status': 'pass' if 'Test passed' in out else 'fail',
                'output': out[:2000]
            }
        except subprocess.CalledProcessError as e:
            return {'test': 'cuda_vectoradd', 'status': 'fail', 'error': str(e), 'output': e.output}
        except Exception as e:
            return {'test': 'cuda_vectoradd', 'status': 'error', 'message': str(e)}
    
    def test_cuda_devquery(self):
        """Run CUDA deviceQuery"""
        try:
            cuda_samples = Path('/usr/local/cuda/samples')
            if not cuda_samples.exists():
                cuda_samples = Path('/opt/cuda/samples')
            
            devq_dir = cuda_samples / '1_Utilities' / 'deviceQuery'
            if not devq_dir.exists():
                return {'test': 'cuda_devquery', 'status': 'skip', 'message': 'deviceQuery sample not found'}
            
            subprocess.check_call(['make'], cwd=devq_dir, timeout=60)
            out = subprocess.check_output(['./deviceQuery'], cwd=devq_dir, text=True, timeout=30)
            
            return {
                'test': 'cuda_devquery',
                'status': 'pass' if 'Result = PASS' in out else 'fail',
                'output': out[:3000]
            }
        except Exception as e:
            return {'test': 'cuda_devquery', 'status': 'error', 'message': str(e)}
    
    def test_vulkan_info(self):
        """Run vulkaninfo"""
        try:
            out = subprocess.check_output(['vulkaninfo'], text=True, timeout=30)
            return {
                'test': 'vulkan_info',
                'status': 'pass' if 'GPU id' in out else 'fail',
                'output': out[:3000]
            }
        except Exception as e:
            return {'test': 'vulkan_info', 'status': 'error', 'message': str(e)}
    
    def test_vkcube(self):
        """Run vkcube (headless)"""
        try:
            # vkcube may need display; try with xvfb or skip
            out = subprocess.check_output(['vkcube', '--duration', '2'], text=True, timeout=30)
            return {'test': 'vkcube', 'status': 'pass', 'output': out[:2000]}
        except FileNotFoundError:
            return {'test': 'vkcube', 'status': 'skip', 'message': 'vkcube not installed'}
        except Exception as e:
            return {'test': 'vkcube', 'status': 'error', 'message': str(e)}
    
    def test_glxgears(self):
        """Run glxgears (brief)"""
        try:
            out = subprocess.check_output(['timeout', '5', 'glxgears'], text=True, timeout=15)
            return {'test': 'glxgears', 'status': 'pass', 'output': out[:1000]}
        except Exception as e:
            return {'test': 'glxgears', 'status': 'error', 'message': str(e)}
    
    def test_ffmpeg_nvenc(self):
        """Test FFmpeg NVENC"""
        try:
            out = subprocess.check_output(['ffmpeg', '-encoders'], text=True, timeout=10)
            has_nvenc = 'h264_nvenc' in out or 'hevc_nvenc' in out
            return {
                'test': 'ffmpeg_nvenc',
                'status': 'pass' if has_nvenc else 'fail',
                'output': out[:2000]
            }
        except Exception as e:
            return {'test': 'ffmpeg_nvenc', 'status': 'error', 'message': str(e)}
    
    def test_pytorch(self):
        """Test PyTorch CUDA"""
        try:
            code = """
import torch
print(f'CUDA available: {torch.cuda.is_available()}')
if torch.cuda.is_available():
    print(f'Device count: {torch.cuda.device_count()}')
    print(f'Device name: {torch.cuda.get_device_name(0)}')
    x = torch.randn(100, 100).cuda()
    y = torch.randn(100, 100).cuda()
    z = torch.mm(x, y)
    print(f'Matmul result shape: {z.shape}')
    print('PyTorch CUDA test: PASS')
else:
    print('PyTorch CUDA test: FAIL - CUDA not available')
"""
            out = subprocess.check_output([sys.executable, '-c', code], text=True, timeout=60)
            return {
                'test': 'pytorch',
                'status': 'pass' if 'PASS' in out else 'fail',
                'output': out
            }
        except Exception as e:
            return {'test': 'pytorch', 'status': 'error', 'message': str(e)}
    
    # ============================================================
    # TEST SUITES
    # ============================================================
    
    def suite_smoke(self):
        """Quick smoke test suite"""
        tests = ['nvidia_smi', 'cuda_devquery', 'cuda_vectoradd']
        return self._run_test_list(tests)
    
    def suite_compute(self):
        """Compute-focused test suite"""
        tests = ['nvidia_smi', 'cuda_devquery', 'cuda_vectoradd', 'pytorch']
        return self._run_test_list(tests)
    
    def suite_graphics(self):
        """Graphics-focused test suite"""
        tests = ['nvidia_smi', 'vulkan_info', 'vkcube', 'glxgears']
        return self._run_test_list(tests)
    
    def suite_video(self):
        """Video encode/decode test suite"""
        tests = ['nvidia_smi', 'ffmpeg_nvenc']
        return self._run_test_list(tests)
    
    def suite_full(self):
        """Full test suite"""
        tests = ['nvidia_smi', 'cuda_devquery', 'cuda_vectoradd', 'vulkan_info', 'vkcube', 'glxgears', 'ffmpeg_nvenc', 'pytorch']
        return self._run_test_list(tests)
    
    def _run_test_list(self, tests):
        """Run a list of tests"""
        results = []
        for test in tests:
            print(f"[{datetime.now()}] Running suite test: {test}")
            result = self._run_test(test, {})
            results.append(result)
            print(f"  -> {result.get('status', 'unknown')}")
        
        passed = sum(1 for r in results if r.get('status') == 'pass')
        failed = sum(1 for r in results if r.get('status') in ('fail', 'error'))
        skipped = sum(1 for r in results if r.get('status') == 'skip')
        
        return {
            'suite': 'custom',
            'total': len(tests),
            'passed': passed,
            'failed': failed,
            'skipped': skipped,
            'results': results
        }

def main():
    parser = argparse.ArgumentParser(description="Guest Test Agent")
    parser.add_argument('--port', type=int, default=9999, help='TCP listen port')
    parser.add_argument('--vsock-cid', type=int, default=3, help='VSOCK CID')
    parser.add_argument('--vsock-port', type=int, default=9999, help='VSOCK port')
    parser.add_argument('--run', help='Run single test and exit (for standalone mode)')
    parser.add_argument('--suite', help='Run test suite and exit')
    args = parser.parse_args()
    
    agent = GuestTestAgent(listen_port=args.port, vsock_cid=args.vsock_cid, vsock_port=args.vsock_port)
    
    if args.run:
        # Standalone mode
        result = agent._run_test(args.run, {})
        print(json.dumps(result, indent=2))
    elif args.suite:
        result = getattr(agent, f'suite_{args.suite}')()
        print(json.dumps(result, indent=2))
    else:
        # Server mode
        try:
            agent.start_server()
        except KeyboardInterrupt:
            print("\nShutting down...")
            agent.running = False

if __name__ == "__main__":
    main()