import os
import sys
import time
import json
import threading
import subprocess
from datetime import datetime
from typing import Dict, List, Optional, Any

import torch


class SystemEnvironment:
    
    @staticmethod
    def get_nvidia_smi_info() -> str:
        try:
            result = subprocess.run(
                ["nvidia-smi"], 
                capture_output=True, 
                text=True, 
                timeout=5
            )
            return result.stdout if result.returncode == 0 else "nvidia-smi returned an error."
        except Exception as e:
            return f"nvidia-smi not available or failed: {e}"

    @staticmethod
    def get_pytorch_info() -> Dict[str, Any]:
        return {
            "python_version": sys.version,
            "torch_version": torch.__version__,
            "cuda_available": torch.cuda.is_available(),
            "cuda_version": torch.version.cuda if torch.cuda.is_available() else None,
            "cudnn_version": torch.backends.cudnn.version() if torch.cuda.is_available() else None,
            "alloc_conf": os.environ.get("PYTORCH_CUDA_ALLOC_CONF", "Not Set"),
        }

    @staticmethod
    def generate_system_report() -> str:
        report = []
        report.append(f"Date: {datetime.utcnow().strftime('%a %b %d %H:%M:%S UTC %Y')}")
        report.append(f"Sys Platform: {sys.platform}")
        report.append("\n=== PyTorch Info ===")
        pt_info = SystemEnvironment.get_pytorch_info()
        for k, v in pt_info.items():
            report.append(f"{k}: {v}")
            
        report.append("\n=== NVIDIA-SMI ===")
        report.append(SystemEnvironment.get_nvidia_smi_info())
        
        return "\n".join(report)


class ResourceTracker:
    
    def __init__(self, polling_interval: float = 0.5, use_cuda_native: bool = True):
        self.polling_interval = polling_interval
        self.use_cuda_native = use_cuda_native
        
        # State
        self.is_running = False
        self._thread: Optional[threading.Thread] = None
        
        # Data
        self.timestamps: List[float] = []
        self.gpu_utilization: List[float] = []
        self.memory_allocated_mb: List[float] = []
        self.memory_reserved_mb: List[float] = []
        
        self._start_time: float = 0.0

    def _poll_metrics(self):
        while self.is_running:
            current_time = time.time() - self._start_time
            self.timestamps.append(current_time)
            
            # 1. PyTorch CUDA Native memory (very exact per-process)
            if torch.cuda.is_available():
                alloc_mb = torch.cuda.memory_allocated() / (1024 ** 2)
                resrv_mb = torch.cuda.memory_reserved() / (1024 ** 2)
            else:
                alloc_mb, resrv_mb = 0.0, 0.0
                
            self.memory_allocated_mb.append(alloc_mb)
            self.memory_reserved_mb.append(resrv_mb)
            
            # 2. System Level GPU Utilization (via fast nvidia-smi call)
            util = 0.0
            if torch.cuda.is_available():
                    # Query utilization
                    res = subprocess.run(
                        ["nvidia-smi", "--query-gpu=utilization.gpu", "--format=csv,noheader,nounits"],
                        capture_output=True, text=True
                    )
                    if res.returncode == 0:
                        lines = res.stdout.strip().split('\n')
                        if lines and lines[0].isdigit():
                            util = float(lines[0])
                except Exception:
                    pass
            self.gpu_utilization.append(util)
            
            time.sleep(self.polling_interval)

    def start(self):
        if self.is_running:
            return
        
        self.is_running = True
        self.timestamps.clear()
        self.gpu_utilization.clear()
        self.memory_allocated_mb.clear()
        self.memory_reserved_mb.clear()
        
        if torch.cuda.is_available():
            torch.cuda.reset_peak_memory_stats()
            
        self._start_time = time.time()
        self._thread = threading.Thread(target=self._poll_metrics, daemon=True)
        self._thread.start()

    def stop(self) -> Dict[str, List[float]]:
        self.is_running = False
        if self._thread is not None:
            self._thread.join(timeout=self.polling_interval * 2)
            
        return self.get_results()
        
    def get_results(self) -> Dict[str, List[float]]:
        return {
            "timestamps_sec": list(self.timestamps),
            "gpu_utilization_pct": list(self.gpu_utilization),
            "memory_allocated_mb": list(self.memory_allocated_mb),
            "memory_reserved_mb": list(self.memory_reserved_mb),
        }
        
    def save_tabular_log(self, filepath: str, label: str = "Workload"):
        with open(filepath, "w") as f:
            f.write(f"=== Resource Tracking: {label} ===\n")
            f.write(f"Polling Interval: {self.polling_interval}s\n")
            f.write("-" * 65 + "\n")
            f.write(f"{'Time(s)':>10} | {'GPU Util(%)':>15} | {'Allocated(MB)':>15} | {'Reserved(MB)':>15}\n")
            f.write("-" * 65 + "\n")
            
            for t, u, a, r in zip(self.timestamps, self.gpu_utilization, self.memory_allocated_mb, self.memory_reserved_mb):
                f.write(f"{t:>10.2f} | {u:>15.1f} | {a:>15.1f} | {r:>15.1f}\n")
            
            f.write("-" * 65 + "\n")
            if self.memory_allocated_mb:
                f.write(f"Peak Allocated: {max(self.memory_allocated_mb):.1f} MB\n")
            if self.memory_reserved_mb:
                f.write(f"Peak Reserved: {max(self.memory_reserved_mb):.1f} MB\n")

if __name__ == "__main__":
    # Smoke test structure
    print(SystemEnvironment.generate_system_report())