"""
VirtueGuardClient - Core safety client for VirtueAI.
Provides synchronous and asynchronous safety checks with MLflow tracing.
"""
import os
import logging
from enum import Enum
from typing import Optional, Dict, Any
from together import Together, AsyncTogether

# --- Optional MLflow Dependency ---
try:
    import mlflow
    def trace_decorator(**kwargs):
        return mlflow.trace(**kwargs)
except ImportError:
    # If MLflow is missing, define a dummy decorator
    def trace_decorator(**kwargs):
        def decorator(func):
            return func
        return decorator

class VirtueAIModel(Enum):
    VIRTUE_GUARD_TEXT_LITE = "Virtue-AI/VirtueGuard-Text-Lite"

class VirtueGuardClient:
    """
    Client for performing VirtueGuard safety checks.
    Can be used directly in Python code or wrapped by Agents.
    """
    def __init__(
        self, 
        api_key: Optional[str] = None,
        safety_model: VirtueAIModel = VirtueAIModel.VIRTUE_GUARD_TEXT_LITE
    ):
        self.api_key = api_key or os.environ.get("VIRTUEAI_API_KEY")
        if not self.api_key:
            raise ValueError(
                "VIRTUEAI_API_KEY not found. Please provide it in init or set VIRTUEAI_API_KEY env var."
            )
            
        self.safety_model = safety_model.value
        self.client = Together(api_key=self.api_key)
        self.async_client = AsyncTogether(api_key=self.api_key)

    @trace_decorator(name="VirtueGuard.evaluate_prompt", span_type="GUARD")
    def evaluate_prompt(self, text: str) -> Dict[str, Any]:
        """
        Evaluate a user prompt for safety.
        Returns a dictionary with status, verdict, and metadata.
        """
        return self._run_check(text, check_type="prompt")

    @trace_decorator(name="VirtueGuard.evaluate_response", span_type="GUARD")
    def evaluate_response(self, text: str) -> Dict[str, Any]:
        """
        Evaluate a model response for safety.
        Returns a dictionary with status, verdict, and metadata.
        """
        return self._run_check(text, check_type="response")
    
    @trace_decorator(name="VirtueGuard.evaluate_prompt_async", span_type="GUARD")
    async def evaluate_prompt_async(self, text: str) -> Dict[str, Any]:
        """Async version of evaluate_prompt."""
        return await self._run_check_async(text, check_type="prompt")

    @trace_decorator(name="VirtueGuard.evaluate_response_async", span_type="GUARD")
    async def evaluate_response_async(self, text: str) -> Dict[str, Any]:
        """Async version of evaluate_response."""
        return await self._run_check_async(text, check_type="response")

    def _run_check(self, text: str, check_type: str) -> Dict[str, Any]:
        try:
            logging.debug(f"Checking {check_type} safety for: {text[:20]}...")
            
            response = self.client.chat.completions.create(
                model=self.safety_model,
                messages=[{"role": "user", "content": text}],
                stream=False,
            )
            return self._parse_response(response, check_type)
        except Exception as e:
            return self._error_response(e, check_type)

    async def _run_check_async(self, text: str, check_type: str) -> Dict[str, Any]:
        try:
            logging.debug(f"Checking {check_type} safety (async) for: {text[:20]}...")
            
            response = await self.async_client.chat.completions.create(
                model=self.safety_model,
                messages=[{"role": "user", "content": text}],
                stream=False,
            )
            return self._parse_response(response, check_type)
        except Exception as e:
            return self._error_response(e, check_type)

    def _parse_response(self, response: Any, check_type: str) -> Dict[str, Any]:
        raw_verdict = response.choices[0].message.content.strip()
        is_safe = raw_verdict.lower().startswith("safe")
        
        return {
            "status": "safe" if is_safe else "unsafe",
            "verdict": raw_verdict,
            "model": self.safety_model,
            "check_type": check_type,
            "action": "ALLOW" if is_safe else "BLOCK"
        }

    def _error_response(self, error: Exception, check_type: str) -> Dict[str, Any]:
        logging.error(f"Error during safety check: {error}")
        return {
            "status": "error",
            "verdict": "undetermined",
            "error": str(error),
            "check_type": check_type,
            "action": "BLOCK" # Fail safe
        }
