"""
VirtueGuardResponsesAgent - Guardrail-protected chat agent using MLflow ResponsesAgent
Compatible with OpenAI Response format and Databricks Model Serving
"""

import json
import logging
import uuid
from typing import Generator, Optional, List

import mlflow
from mlflow.entities import SpanType
from mlflow.pyfunc import ResponsesAgent
from mlflow.types.responses import (
    ResponsesAgentRequest,
    ResponsesAgentResponse,
    ResponsesAgentStreamEvent,
    to_chat_completions_input,
)

from virtueai.guard.client import VirtueGuardClient, VirtueAIModel

class VirtueGuardResponsesAgent(ResponsesAgent):
    """
    VirtueGuard implementation as a Databricks ResponsesAgent.
    Acts as a functional Guard that accepts text/messages and returns a 
    structured safety assessment.
    """
    
    def __init__(
        self,
        safety_model: VirtueAIModel = VirtueAIModel.VIRTUE_GUARD_TEXT_LITE,
    ):
        self.safety_model = safety_model
        self.client: Optional[VirtueGuardClient] = None

    def _get_client(self) -> VirtueGuardClient:
        """
        Lazily initialize the VirtueGuardClient.
        This ensures the API key is read from the serving environment at runtime.
        """
        if self.client is None:
            # VirtueGuardClient handles env var validation internally
            self.client = VirtueGuardClient(safety_model=self.safety_model)
        return self.client

    def predict(self, request: ResponsesAgentRequest) -> ResponsesAgentResponse:
        """
        Evaluate inputs for safety and return a structured assessment.
        """
        client = self._get_client()
        
        # Convert request input to standard chat completion format using MLflow helper
        messages = to_chat_completions_input(request.input)
        
        # Extract user content for safety check
        user_content = " ".join(m["content"] for m in messages if m["role"] == "user")
        
        # Use the client to evaluate
        # Tracing is handled automatically by the client's decorator
        assessment = client.evaluate_prompt(user_content)
        
        output_item = self.create_text_output_item(
            text=json.dumps(assessment),
            id=str(uuid.uuid4())
        )

        return ResponsesAgentResponse(
            output=[output_item],
            custom_outputs={
                "safety_status": assessment["status"],
                "verdict": assessment["verdict"],
                "action": assessment.get("action", "BLOCK")
            }
        )
    
    def predict_stream(self, request: ResponsesAgentRequest) -> Generator[ResponsesAgentStreamEvent, None, None]:
        """
        Streaming interface - delegates to predict() as safety checks are atomic.
        """
        response = self.predict(request)
        yield ResponsesAgentStreamEvent(
            item=response.output[0],
            type="response.output_item.added"
        )

if __name__ == "__main__":
    import os
    # For local testing only
    AGENT = VirtueGuardResponsesAgent(
        safety_model=VirtueAIModel.VIRTUE_GUARD_TEXT_LITE
    )
    mlflow.models.set_model(AGENT)
