from __future__ import annotations

import json
from datetime import datetime
from typing import Any, Callable, Dict, List, Optional, TypedDict, Union
from aryaxai.common.xai_uris import RUN_GUARDRAILS_URI , RUN_GUARDRAILS_PARALLEL_URI
from aryaxai.core.project import Project
from opentelemetry import trace
import time
import asyncio
import requests
from .guard_template import Guard
from dataclasses import dataclass

from agents import (
    Agent,
    RunContextWrapper,
    TResponseInputItem,
    input_guardrail,
    output_guardrail,
    ModelSettings,
    ModelTracing,
)

@dataclass
class GuardrailFunctionOutput:
    """The output of a guardrail function."""

    output_info: Any
    """
    Optional information about the guardrail's output. For example, the guardrail could include
    information about the checks it performed and granular results.
    """

    tripwire_triggered: bool
    """
    Whether the tripwire was triggered. If triggered, the agent's execution will be halted.
    """

    sanitized_content: str

class GuardrailRunResult(TypedDict, total=False):
    success: bool
    details: Dict[str, Any]
    validated_output: Any
    validation_passed: bool
    sanitized_output: Any
    duration: float
    latency: str
    on_fail_action: str
    retry_count: int
    max_retries: int
    start_time: str
    end_time: str

class OpenAIAgentsGuardrail:
    """
    Utility for creating OpenAI Agents guardrails that call Guardrails HTTP APIs.
    
    Supports different actions:
    - "block": raises InputGuardrailTripwireTriggered/OutputGuardrailTripwireTriggered
    - "retry": attempts to sanitize content using LLM and passes sanitized content to handoff agent
    - "warn": logs the issue but allows content to pass through
    """

    def __init__(
        self,
        project: Optional[Project],
        model: Optional[Any] = None,
    ) -> None:
        if project is not None:
            self.client = project.api_client
            self.project_name = project.project_name
            
        self.logs: List[Dict[str, Any]] = []
        self.max_retries = 2
        self.retry_delay = 1.0 
        self.tracer = trace.get_tracer(__name__)
        self.model = model

    def create_input_guardrail(
        self,
        guards: Union[List[str], List[Dict[str, Any]], str, Dict[str, Any]],
        action: str = "block",
        name: str = "input_guardrail"
    ) -> Callable:
        """
        Create an input guardrail function for OpenAI Agents.
        
        Args:
            guards: List of guard specifications or single guard
            action: 'block' | 'retry' | 'warn'
            name: Name for the guardrail function
        """
        if isinstance(guards, (str, dict)):
            guards = [guards]

        @input_guardrail
        async def guardrail_function(
            ctx: RunContextWrapper[None], 
            agent: Agent, 
            input: str | list[TResponseInputItem]
        ) -> GuardrailFunctionOutput:
            # Convert input to string for processing
            if isinstance(input, list):
                # Handle list of input items (messages)
                input_text = ""
                for item in input:
                    if hasattr(item, 'content'):
                        input_text += str(item.content) + " "
                    else:
                        input_text += str(item) + " "
                input_text = input_text.strip()
            else:
                input_text = str(input)

            # Process through all guards in parallel
            current_content, tripwire_triggered, output_info = await self._apply_guardrail_parallel(
                content=input_text,
                guards=guards,
                guardrail_type="input",
                action=action,
                agent_name=agent.name,
            )

            return GuardrailFunctionOutput(
                output_info=output_info,
                tripwire_triggered=tripwire_triggered,
                sanitized_content=current_content if action == "retry" else input_text
            )

        # Set function name for debugging
        guardrail_function.__name__ = name
        return guardrail_function

    def create_output_guardrail(
        self,
        guards: Union[List[str], List[Dict[str, Any]], str, Dict[str, Any]],
        action: str = "block",
        name: str = "output_guardrail"
    ) -> Callable:
        """
        Create an output guardrail function for OpenAI Agents.
        
        Args:
            guards: List of guard specifications or single guard
            action: 'block' | 'retry' | 'warn'
            name: Name for the guardrail function
        """
        if isinstance(guards, (str, dict)):
            guards = [guards]

        @output_guardrail
        async def guardrail_function(
            ctx: RunContextWrapper, 
            agent: Agent, 
            output: Any
        ) -> GuardrailFunctionOutput:
            # Extract text content from output
            if hasattr(output, 'response'):
                output_text = str(output.response)
            elif hasattr(output, 'content'):
                output_text = str(output.content)
            else:
                output_text = str(output)

            # Process through all guards in parallel
            current_content, tripwire_triggered, output_info = await self._apply_guardrail_parallel(
                content=output_text,
                guards=guards,
                guardrail_type="output",
                action=action,
                agent_name=agent.name,
            )

            return GuardrailFunctionOutput(
                output_info=output_info,
                tripwire_triggered=tripwire_triggered,
                sanitized_content=current_content if action == "retry" else output_text
            )

        # Set function name for debugging
        guardrail_function.__name__ = name
        return guardrail_function

    async def _apply_guardrail_parallel(
        self,
        content: Any,
        guards: List[Union[str, Dict[str, Any]]],
        guardrail_type: str,
        action: str,
        agent_name: str,
    ) -> tuple[Any, bool, Dict[str, Any]]:
        """
        Run multiple guardrails in parallel using batch endpoint.
        
        Returns:
            tuple: (processed_content, tripwire_triggered, output_info)
        """
        current_content = content
        tripwire_triggered = False
        output_info = {}
        retry_count = 0

        try:
            parent_span = trace.get_current_span()
            if parent_span is not None:
                ctx = trace.set_span_in_context(parent_span)
                with self.tracer.start_as_current_span(f"guardrails:{guardrail_type}", context=ctx) as parent_gr_span:
                    parent_gr_span.set_attribute("component", str(agent_name))
                    parent_gr_span.set_attribute("content_type", guardrail_type)

                    while retry_count <= self.max_retries:
                        # Prepare payload for parallel guardrail execution
                        guard_specs = [guard if isinstance(guard, dict) else {"name": guard} for guard in guards]
                        payload = {
                            "input_data": current_content,
                            "guards": guard_specs
                        }

                        # Call parallel guardrail endpoint
                        start_time = datetime.now()
                        response = self.client.post(
                            RUN_GUARDRAILS_PARALLEL_URI,
                            payload=payload,
                        )
                        end_time = datetime.now()
                        parallel_result = response

                        # Add timing information
                        parallel_result.update({
                            "start_time": start_time.isoformat(),
                            "end_time": end_time.isoformat(),
                            "duration": (end_time - start_time).total_seconds()
                        })

                        parent_gr_span.set_attribute("start_time", str(parallel_result.get("start_time", "")))
                        parent_gr_span.set_attribute("end_time", str(parallel_result.get("end_time", "")))
                        parent_gr_span.set_attribute("duration", float(parallel_result.get("duration", 0.0)))

                        if not parallel_result.get("success", False):
                            output_info["parallel_execution_error"] = parallel_result.get("details", {})
                            return current_content, tripwire_triggered, output_info

                        # Process each guardrail result
                        detected_issue = False
                        for guard_result in parallel_result.get("details", []):
                            guard_name = guard_result.get("name", "unknown")
                            run_result: GuardrailRunResult = {
                                "success": parallel_result.get("success"),
                                "details": guard_result,
                                "validated_output": guard_result.get("validated_output"),
                                "validation_passed": guard_result.get("validation_passed", False),
                                "sanitized_output": guard_result.get("sanitized_output", current_content),
                                "duration": guard_result.get("duration", 0.0),
                                "latency": guard_result.get("latency", "0 ms"),
                                "start_time": parallel_result.get("start_time", ""),
                                "end_time": parallel_result.get("end_time", ""),
                                "retry_count": retry_count,
                                "max_retries": self.max_retries
                            }
                            run_result["response"] = guard_result
                            run_result["input"] = current_content

                            # Log and handle each guard result
                            current_content, is_triggered = await self._handle_action(
                                original=current_content,
                                run_result=run_result,
                                action=f"retry_{retry_count}" if retry_count > 0 else action,
                                agent_name=agent_name,
                                guardrail_type=guardrail_type,
                                guard_name=guard_name,
                                parent_span=parent_gr_span
                            )

                            if is_triggered:
                                detected_issue = True
                                tripwire_triggered = True
                                output_info[f"guard_{guard_name}"] = run_result

                        if detected_issue and action == "retry" and self.model is not None and retry_count < self.max_retries:
                            # Sanitize content using LLM
                            prompt = self._build_sanitize_prompt("combined", current_content, guardrail_type)
                            try:
                                sanitized = await self._invoke_llm(prompt)
                                current_content = sanitized
                            except Exception:
                                pass  # Keep original content if sanitization fails
                            retry_count += 1
                            await self._async_sleep(self.retry_delay)
                            continue
                        else:
                            # No issues or no retries left
                            output_info["retry_count"] = retry_count
                            output_info["final_content"] = current_content
                            return current_content, tripwire_triggered, output_info

            # Fallback if no parent span
            output_info["retry_count"] = retry_count
            output_info["final_content"] = current_content
            return current_content, tripwire_triggered, output_info

        except Exception as e:
            output_info["error"] = f"Parallel guardrail execution failed: {str(e)}"
            return current_content, tripwire_triggered, output_info

    async def _handle_action(
        self,
        original: Any,
        run_result: GuardrailRunResult,
        action: str,
        agent_name: str,
        guardrail_type: str,
        guard_name: str,
        parent_span: Optional[Any]
    ) -> tuple[Any, bool]:
        """
        Handle the action for a single guardrail result.
        
        Returns:
            tuple: (processed_content, is_triggered)
        """
        validation_passed = bool(run_result.get("validation_passed", True))
        detected_issue = not validation_passed or not run_result.get("success", True)

        if parent_span is not None:
            try:
                with self.tracer.start_as_current_span(
                    f"guard:{guard_name}",
                    context=trace.set_span_in_context(parent_span)
                ) as gr_span:
                    gr_span.set_attribute("component", str(agent_name))
                    gr_span.set_attribute("guard", str(guard_name))
                    gr_span.set_attribute("content_type", guardrail_type)
                    gr_span.set_attribute("detected", detected_issue)
                    gr_span.set_attribute("action", action)
                    gr_span.set_attribute("input.value", self._safe_str(run_result.get("input")))
                    gr_span.set_attribute("output.value", json.dumps(run_result.get("response")))
                    gr_span.set_attribute("start_time", str(run_result.get("start_time", "")))
                    gr_span.set_attribute("end_time", str(run_result.get("end_time", "")))
                    gr_span.set_attribute("duration", float(run_result.get("duration", 0.0)))
            except Exception:
                pass

        # Log guardrail result without creating a new span
        self._log_guardrail_result(
            run_result=run_result,
            action=action,
            agent_name=agent_name,
            guardrail_type=guardrail_type,
            guard_name=guard_name,
        )

        if detected_issue:
            if action == "block":
                return original, True
            elif "retry" in action:
                return run_result.get("sanitized_output", original), True
            else:  # warn
                return original, False
        return original, False

    async def _call_run_guardrail(
        self, 
        input_data: Any, 
        guard: Dict[str, Any], 
        guardrail_type: str
    ) -> GuardrailRunResult:
        """Call the guardrails HTTP API"""
        uri = RUN_GUARDRAILS_URI
        input_text = str(input_data)
        
        start_time = datetime.now()
        try:
            body = {"input_data": input_text, "guard": guard}
            data = self.client.post(uri, body)
            
            end_time = datetime.now()
            
            details = data.get("details", {}) if isinstance(data, dict) else {}
            result: GuardrailRunResult = {
                "success": bool(data.get("success", False)) if isinstance(data, dict) else False,
                "details": details if isinstance(details, dict) else {},
                "start_time": start_time.isoformat(),
                "end_time": end_time.isoformat(),
            }
            
            if "duration" not in details:
                result["duration"] = (end_time - start_time).total_seconds()
                
            if isinstance(details, dict):
                if "validated_output" in details:
                    result["validated_output"] = details["validated_output"]
                if "validation_passed" in details:
                    result["validation_passed"] = details["validation_passed"]
                if "sanitized_output" in details:
                    result["sanitized_output"] = details["sanitized_output"]
                if "duration" in details:
                    result["duration"] = details["duration"]
                if "latency" in details:
                    result["latency"] = details["latency"]
                
            result["retry_count"] = 0
            result["max_retries"] = self.max_retries
            result["response"] = data
            result["input"] = input_text

            return result
            
        except Exception as exc:
            end_time = datetime.now()
            raise exc

    def _log_guardrail_result(
        self,
        run_result: GuardrailRunResult,
        action: str,
        agent_name: str,
        guardrail_type: str,
        guard_name: str,
    ) -> None:
        """Log guardrail results without creating a new span"""
        validation_passed = bool(run_result.get("validation_passed", True))
        detected_issue = not validation_passed or not run_result.get("success", True)
        
        # Append to logs list instead of creating a new span
        log_entry = {
            "guard_name": guard_name,
            "guardrail_type": guardrail_type,
            "agent_name": agent_name,
            "action": action,
            "detected_issue": detected_issue,
            "start_time": run_result.get("start_time", ""),
            "end_time": run_result.get("end_time", ""),
            "duration": float(run_result.get("duration", 0.0)),
            "input": self._safe_str(run_result.get("input")),
            "output": self._safe_str(run_result.get("response"))
        }
        self.logs.append(log_entry)

    def _build_sanitize_prompt(self, guard_name: str, content: Any, guardrail_type: str) -> str:
        """Build a prompt for the LLM to sanitize the content according to the guardrail type"""
        instructions = {
            "Detect PII": "Sanitize the following text by removing or masking any personally identifiable information (PII). Do not change anything else.",
            "NSFW Text": "Sanitize the following text by removing or masking any not safe for work (NSFW) content. Do not change anything else.",
            "Ban List": "Sanitize the following text by removing or masking any banned words. Do not change anything else.",
            "Bias Check": "Sanitize the following text by removing or masking any biased language. Do not change anything else.",
            "Competitor Check": "Sanitize the following text by removing or masking any competitor names. Do not change anything else.",
            "Correct Language": "Sanitize the following text by correcting the language to the expected language. Do not change anything else.",
            "Gibberish Text": "Sanitize the following text by removing or correcting any gibberish. Do not change anything else.",
            "Profanity Free": "Sanitize the following text by removing or masking any profanity. Do not change anything else.",
            "Secrets Present": "Sanitize the following text by removing or masking any secrets. Do not change anything else.",
            "Toxic Language": "Sanitize the following text by removing or masking any toxic language. Do not change anything else.",
            "Contains String": "Sanitize the following text by removing or masking the specified substring. Do not change anything else.",
            "Detect Jailbreak": "Sanitize the following text by removing or masking any jailbreak attempts. Do not change anything else.",
            "Endpoint Is Reachable": "Sanitize the following text by ensuring any mentioned endpoints are reachable. Do not change anything else.",
            "Ends With": "Sanitize the following text by ensuring it ends with the specified string. Do not change anything else.",
            "Has Url": "Sanitize the following text by removing or masking any URLs. Do not change anything else.",
            "Lower Case": "Sanitize the following text by converting it to lower case. Do not change anything else.",
            "Mentions Drugs": "Sanitize the following text by removing or masking any mentions of drugs. Do not change anything else.",
            "One Line": "Sanitize the following text by ensuring it is a single line. Do not change anything else.",
            "Reading Time": "Sanitize the following text by ensuring its reading time matches the specified value. Do not change anything else.",
            "Redundant Sentences": "Sanitize the following text by removing redundant sentences. Do not change anything else.",
            "Regex Match": "Sanitize the following text by ensuring it matches the specified regex. Do not change anything else.",
            "Sql Column Presence": "Sanitize the following text by ensuring specified SQL columns are present. Do not change anything else.",
            "Two Words": "Sanitize the following text by ensuring it contains only two words. Do not change anything else.",
            "Upper Case": "Sanitize the following text by converting it to upper case. Do not change anything else.",
            "Valid Choices": "Sanitize the following text by ensuring it matches one of the valid choices. Do not change anything else.",
            "Valid Json": "Sanitize the following text by ensuring it is valid JSON. Do not change anything else.",
            "Valid Length": "Sanitize the following text by ensuring its length is valid. Do not change anything else.",
            "Valid Range": "Sanitize the following text by ensuring its value is within the valid range. Do not change anything else.",
            "Valid URL": "Sanitize the following text by ensuring it is a valid URL. Do not change anything else.",
            "Web Sanitization": "Sanitize the following text by removing any unsafe web content. Do not change anything else.",
        }
        
        instruction = instructions.get(guard_name, "Sanitize the following text according to the guardrail requirements. Do not change anything else.")
        prompt = f"{instruction}\n\nContent:\n{content}"
        return prompt

    async def _invoke_llm(self, prompt: str) -> str:
        """Invoke the LLM for content sanitization"""
        if self.model is None:
            return prompt  # Return original if no LLM available
            
        try:
            data = await self.model.get_response(
                system_instructions="Based on the input you have to provide the best and accurate results",
                input=prompt,
                model_settings=ModelSettings(temperature=0.1),
                tools=[],
                output_schema=None,
                handoffs=[],
                tracing=ModelTracing.DISABLED,
                previous_response_id=None
            )
            return str(data.output[0].content[0].text)
        except Exception:
            return prompt  # Return original on error

    async def _async_sleep(self, seconds: float) -> None:
        """Async sleep utility"""
        await asyncio.sleep(seconds)

    @staticmethod
    def _safe_str(value: Any) -> str:
        """Safely convert any value to string for logging"""
        try:
            if isinstance(value, (str, int, float, bool)) or value is None:
                return str(value)
            if hasattr(value, "content"):
                return str(getattr(value, "content", ""))
            
            if isinstance(value, (list, tuple)):
                parts = []
                for item in value:
                    parts.append(OpenAIAgentsGuardrail._safe_str(item))
                return ", ".join(parts)
            
            if isinstance(value, dict):
                safe_dict: Dict[str, Any] = {}
                for k, v in value.items():
                    key = str(k)
                    if isinstance(v, (str, int, float, bool)) or v is None:
                        safe_dict[key] = v
                    elif hasattr(v, "content"):
                        safe_dict[key] = str(getattr(v, "content", ""))
                    else:
                        safe_dict[key] = str(v)
                return json.dumps(safe_dict, ensure_ascii=False)
            
            return str(value)
        except Exception:
            return "<unserializable>"

def create_guardrail(project: Project, model: Optional[Any] = None) -> OpenAIAgentsGuardrail:
    """Quick factory function to create a guardrail instance with a project"""
    return OpenAIAgentsGuardrail(project=project, model=model)