import os
import json
import time
import threading
import subprocess
from datetime import datetime
from abc import ABC, abstractmethod


class BaseAgent(ABC):
    def __init__(self, name, status_file="agents/agent_status.json"):
        self.name = name
        self.status_file = status_file
        self.status = "idle"
        self.progress = 0
        self.message = ""
        self.start_time = None
        self.gpu_available = self._check_gpu()
        self._lock = threading.Lock()
        self._running = False
        os.makedirs(os.path.dirname(status_file), exist_ok=True)
        self._write_status()

    def _check_gpu(self):
        try:
            result = subprocess.run(
                ["nvidia-smi", "--query-gpu=name,memory.total", "--format=csv,noheader"],
                capture_output=True, text=True, timeout=5
            )
            if result.returncode == 0 and result.stdout.strip():
                return result.stdout.strip()
        except Exception:
            pass
        return None

    def get_gpu_stats(self):
        try:
            result = subprocess.run(
                ["nvidia-smi", "--query-gpu=temperature.gpu,memory.used,memory.total,utilization.gpu",
                 "--format=csv,noheader,nounits"],
                capture_output=True, text=True, timeout=5
            )
            if result.returncode == 0:
                parts = result.stdout.strip().split(",")
                return {
                    "temp": int(parts[0].strip()),
                    "mem_used": int(parts[1].strip()),
                    "mem_total": int(parts[2].strip()),
                    "util": int(parts[3].strip())
                }
        except Exception:
            pass
        return None

    def set_status(self, status, progress=None, message=""):
        with self._lock:
            self.status = status
            if progress is not None:
                self.progress = progress
            if message:
                self.message = message
            if status == "running" and self.start_time is None:
                self.start_time = time.time()
            elif status in ("completed", "failed", "idle"):
                self.start_time = None
            self._write_status()

    def _write_status(self):
        try:
            all_status = {}
            if os.path.exists(self.status_file):
                with open(self.status_file, "r") as f:
                    all_status = json.load(f)

            gpu = self.get_gpu_stats()
            all_status[self.name] = {
                "status": self.status,
                "progress": self.progress,
                "message": self.message,
                "gpu_available": self.gpu_available is not None,
                "gpu_stats": gpu,
                "start_time": self.start_time,
                "updated_at": datetime.now().isoformat()
            }

            with open(self.status_file, "w") as f:
                json.dump(all_status, f, indent=2)
        except Exception:
            pass

    def log(self, msg):
        ts = datetime.now().strftime("%H:%M:%S")
        print(f"[{ts}] [{self.name}] {msg}")
        self.message = msg
        self._write_status()

    def wait_for_gpu(self, max_temp=80, max_mem_pct=90):
        self.log(f"Waiting for GPU (max temp: {max_temp}C, max mem: {max_mem_pct}%)...")
        while True:
            stats = self.get_gpu_stats()
            if stats is None:
                self.log("No GPU stats available, proceeding anyway")
                return True
            temp_ok = stats["temp"] < max_temp
            mem_pct = (stats["mem_used"] / max(1, stats["mem_total"])) * 100
            mem_ok = mem_pct < max_mem_pct
            if temp_ok and mem_ok:
                self.log(f"GPU ready: {stats['temp']}C, {mem_pct:.0f}% mem used")
                return True
            self.log(f"GPU busy: {stats['temp']}C, {mem_pct:.0f}% mem, {stats['util']}% util - waiting 10s")
            time.sleep(10)

    @abstractmethod
    def run(self, **kwargs):
        pass

    def start(self, **kwargs):
        self._running = True
        self.set_status("running", 0, "Starting...")
        try:
            self.run(**kwargs)
            self.set_status("completed", 100, "Done")
        except Exception as e:
            self.set_status("failed", self.progress, f"Error: {str(e)}")
            self.log(f"FAILED: {e}")
            raise
        finally:
            self._running = False

    def stop(self):
        self._running = False
        self.set_status("stopped", self.progress, "Stopped by user")