import requests
import json
import psutil
try:
    import pynvml
except ImportError:
    pynvml = None

from modules.config import Config

class OllamaManager:
    """
    Ollama API 세션 관리 및 메모리 누수 방지 래퍼
    - GPU VRAM 사용량 모니터링
    - keep_alive 파라미터를 통한 메모리 관리
    """
    def __init__(self, model_name=None, base_url=None):
        self.model_name = model_name or Config.OLLAMA_MODEL
        self.base_url = base_url or Config.OLLAMA_API_URL
        self.generate_url = f"{self.base_url}/api/generate"
        
        self.gpu_available = False
        try:
            if pynvml:
                pynvml.nvmlInit()
                self.handle = pynvml.nvmlDeviceGetHandleByIndex(0) # 0번 GPU (3070 Ti)
                self.gpu_available = True
                print("✅ [OllamaManager] NVIDIA GPU Monitoring On")
            else:
                print("⚠️ [OllamaManager] 'nvidia-ml-py' not installed. GPU monitoring disabled.")
        except Exception as e:
            print(f"⚠️ [OllamaManager] GPU Init Failed: {e}")

    def check_vram(self):
        """현재 GPU VRAM 사용량(GB) 반환"""
        if not self.gpu_available:
            return 0.0
        try:
            info = pynvml.nvmlDeviceGetMemoryInfo(self.handle)
            used_gb = info.used / 1024**3
            return used_gb
        except Exception:
            return 0.0

    def get_gpu_status(self):
        """GPU 종합 상태 반환 (온도, 메모리, 사용률, 이름)"""
        if not self.gpu_available:
            return {"name": "N/A", "temp": 0, "vram_used": 0, "vram_total": 0, "load": 0}
            
        try:
            # GPU 이름
            name = pynvml.nvmlDeviceGetName(self.handle)
            if isinstance(name, bytes):
                name = name.decode('utf-8')
            
            # 온도
            temp = pynvml.nvmlDeviceGetTemperature(self.handle, pynvml.NVML_TEMPERATURE_GPU)
            # 메모리
            mem_info = pynvml.nvmlDeviceGetMemoryInfo(self.handle)
            vram_used = mem_info.used / 1024**3
            vram_total = mem_info.total / 1024**3
            # 사용률
            util = pynvml.nvmlDeviceGetUtilizationRates(self.handle)
            load = util.gpu
            
            return {
                "name": name,
                "temp": temp,
                "vram_used": round(vram_used, 1),
                "vram_total": round(vram_total, 1),
                "load": load
            }
        except Exception as e:
            print(f"⚠️ GPU Status Check Failed: {e}")
            return {"name": "N/A", "temp": 0, "vram_used": 0, "vram_total": 0, "load": 0}

    def request_inference(self, prompt, context_data=None):
        """
        Ollama에 추론 요청
        :param prompt: 시스템 프롬프트 + 사용자 입력
        :param context_data: (Optional) 이전 대화 컨텍스트
        """
        # [5070Ti 최적화] VRAM이 14GB 이상이면 모델 언로드 시도 (16GB 중 여유분 확보)
        vram = self.check_vram()
        if vram > 14.0:
            print(f"⚠️ [OllamaManager] High VRAM Usage ({vram:.1f}GB). Requesting unload.")
            try:
                # keep_alive=0으로 설정하여 모델 즉시 언로드
                requests.post(self.generate_url, 
                              json={"model": self.model_name, "keep_alive": 0}, timeout=5)
            except Exception as e:
                print(f"Warning: Failed to unload model: {e}")

        payload = {
            "model": self.model_name,
            "prompt": prompt,
            "stream": False,
            "format": "json", # JSON 강제
            "options": {
                "num_ctx": 8192,    # [5070Ti 최적화] 컨텍스트 크기 2배 증가 (4096 -> 8192)
                "temperature": 0.2,  # 분석 일관성 유지
                "num_gpu": 1,        # GPU 사용 명시
                "num_thread": 8      # CPU 스레드 수 (9800X3D 활용)
            },
            "keep_alive": "10m"      # [5070Ti 최적화] 10분간 유지 (메모리 여유 있음)
        }
        
        try:
            response = requests.post(self.generate_url, json=payload, timeout=180)  # 타임아웃 증가
            response.raise_for_status()
            return response.json().get('response')
        except Exception as e:
            print(f"❌ Inference Error: {e}")
            return None