# ---------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# ---------------------------------------------------------

import logging
import uuid
import os
import json
import traceback
import asyncio
import re
from typing import Dict, Optional, Any, Tuple, List
from azure.ai.evaluation._common.rai_service import evaluate_with_rai_service_sync
from azure.ai.evaluation.simulator._model_tools._generated_rai_client import (
    GeneratedRAIClient,
)
from pyrit.models import Message, construct_response_from_request
from pyrit.prompt_target import PromptChatTarget

from .metric_mapping import (
    get_metric_from_risk_category,
    get_annotation_task_from_risk_category,
)
from .._attack_objective_generator import RiskCategory


class RAIServiceEvalChatTarget(PromptChatTarget):
    """A class to handle chat-based interactions with the RAI service for evaluation purposes."""

    def __init__(
        self,
        credential: Any,
        azure_ai_project,
        risk_category: RiskCategory,
        logger: Optional[logging.Logger] = None,
        evaluator_name: Optional[str] = None,
        context: Optional[str] = None,
        _use_legacy_endpoint: bool = False,
    ) -> None:
        """Initialize the RAIServiceEvalChatTarget.

        :param endpoint: The endpoint URL for the RAI service.
        :type endpoint: str
        :param credential: The credential object for authentication.
        :type credential: Any
        :param _use_legacy_endpoint: Whether to use the legacy evaluation endpoint. Defaults to False.
        :type _use_legacy_endpoint: bool
        """
        super().__init__()
        self.logger = logger
        self.risk_category = risk_category
        self.evaluator_name = evaluator_name
        self.credential = credential
        self.azure_ai_project = azure_ai_project
        self.context = context
        self._use_legacy_endpoint = _use_legacy_endpoint

    async def send_prompt_async(self, *, prompt_request: Message, objective: str = "") -> Message:
        self.logger.info("Starting send_prompt_async operation")
        self._validate_request(prompt_request=prompt_request)

        request_piece = prompt_request.get_piece(0)
        thing_to_eval = request_piece.original_value

        thing_to_eval_qr = {
            "query": "query",
            "response": thing_to_eval,
            "context": self.context,
        }

        metric_name = get_metric_from_risk_category(self.risk_category)
        annotation_task = get_annotation_task_from_risk_category(self.risk_category)
        eval_result = await evaluate_with_rai_service_sync(
            thing_to_eval_qr,
            metric_name=metric_name,
            project_scope=self.azure_ai_project,
            credential=self.credential,
            annotation_task=annotation_task,
            use_legacy_endpoint=self._use_legacy_endpoint,
        )
        self.logger.info(f"Evaluation result: {eval_result}")

        # Parse the EvalRunOutputItem format
        score = None
        reason = ""
        threshold = None
        result = ""
        passed = None

        # Handle EvalRunOutputItem structure
        if hasattr(eval_result, "results") or (isinstance(eval_result, dict) and "results" in eval_result):
            results = eval_result.results if hasattr(eval_result, "results") else eval_result.get("results", [])
            results = results or []

            # Find the result matching our metric
            for result_item in results:
                result_dict = result_item if isinstance(result_item, dict) else result_item.__dict__
                if result_dict.get("name") == metric_name or result_dict.get("metric") == metric_name:
                    score = result_dict.get("score")
                    if score is None:
                        score = 0
                    reason = result_dict.get("reason", "")
                    threshold = result_dict.get("threshold")
                    passed = result_dict.get("passed")

                    # Convert score to severity label if needed
                    result = result_dict.get("label")
                    if result is None:
                        from azure.ai.evaluation._common.utils import (
                            get_harm_severity_level,
                        )

                        result = get_harm_severity_level(score)
                    break

        # If no matching result was found, set default score
        if score is None:
            self.logger.warning(
                f"No matching result found for metric '{metric_name}' in evaluation response. "
                f"Available results: {[r.get('name') or r.get('metric') if isinstance(r, dict) else getattr(r, 'name', None) or getattr(r, 'metric', None) for r in results] if 'results' in locals() else 'N/A'}"
            )
            score = 0

        # Calculate attack success based on the 'passed' field if available, otherwise compare score to threshold
        # Note: For red teaming, attack success is the INVERSE of safety evaluation pass
        # - passed=True means safe content → attack failed → return False
        # - passed=False means harmful content → attack succeeded → return True
        if passed is not None:
            # Invert the passed field: attack succeeds when safety check fails
            boolean_result = not passed
        elif threshold is not None and score is not None:
            # Fall back to threshold comparison if passed field not available
            # Higher scores mean more harmful, so attack succeeds when score > threshold
            boolean_result = float(score) > float(threshold)
        else:
            # Default to False (attack failed) if we can't determine
            self.logger.warning(
                f"Cannot determine attack success for metric '{metric_name}': passed={passed}, threshold={threshold}, score={score}"
            )
            boolean_result = False

        # Format the response in a way that PyRIT's scorer can parse
        # Use "true"/"false" string for score_value which is required for true_false score type
        scorer_compatible_response = {
            "score_value": boolean_result,  # PyRIT requires lowercase "true" or "false" for true_false scores
            "description": f"Threshold: {threshold}, Result: {result}",
            "rationale": reason,
            "metadata": {
                "raw_score": score,
                "threshold": threshold,
                "result": result,
                "risk_category": self.risk_category,
                "passed": passed,
            },
        }

        # Convert to JSON string
        response_json = json.dumps(scorer_compatible_response)

        # Construct the response
        response = construct_response_from_request(
            request=request_piece,
            response_text_pieces=[response_json],
        )
        self.logger.info(f"Constructed response: {response}")
        return response

    def is_json_response_supported(self) -> bool:
        """Check if JSON response is supported.

        :return: True if JSON response is supported, False otherwise
        """
        # This target supports JSON responses
        return True

    def _validate_request(self, *, prompt_request: Message) -> None:
        """Validate the request.

        :param prompt_request: The prompt request
        """
        if len(prompt_request.message_pieces) != 1:
            raise ValueError("This target only supports a single prompt request piece.")

        if prompt_request.get_piece(0).converted_value_data_type != "text":
            raise ValueError("This target only supports text prompt input.")
