"""
OpenAI Agents SDK Adapter - Wrapper for the OpenAI Agents SDK.

This adapter provides a unified interface for using the OpenAI Agents SDK
which offers handoffs, sessions, guardrails, and built-in usage reporting.

Usage:
    from agent_orchestrator.adapters.openai_agents import OpenAIAgentsAdapter

    adapter = OpenAIAgentsAdapter(
        agent_id="openai-coder",
        api_key=os.environ["OPENAI_API_KEY"],
    )
    response = await adapter.execute("Write unit tests", context)
"""

import asyncio
from dataclasses import dataclass
from datetime import datetime
from typing import Any, AsyncIterator, Optional

from .base import (
    LLMAdapter,
    AgentResponse,
    AgentStatus,
    UsageStats,
    PromptBuilder,
)
from .pricing import estimate_cost
from ..journal.status_packet import (
    StatusPacket,
    TaskArtifacts,
    create_success_packet,
    create_failed_packet,
)
from ..secrets.redactor import SecretRedactor


@dataclass
class OpenAIAgentsConfig:
    """Configuration for OpenAI Agents adapter."""

    api_key: str
    model: str = "gpt-4o"
    max_tokens: int = 4096
    temperature: float = 0.7
    timeout_seconds: int = 300


class OpenAIAgentsAdapter(LLMAdapter):
    """
    Adapter for the OpenAI Agents SDK.

    Provides access to OpenAI's agent capabilities including:
    - Handoffs between specialized agents
    - Session management
    - Guardrails for input validation
    - Built-in usage tracking

    Note: Requires the openai-agents package to be installed:
        pip install openai-agents
    """

    def __init__(
        self,
        agent_id: str,
        api_key: str,
        model: str = "gpt-4o",
        config: Optional[OpenAIAgentsConfig] = None,
    ) -> None:
        """
        Initialize the OpenAI Agents adapter.

        Args:
            agent_id: Unique identifier for this agent instance
            api_key: OpenAI API key
            model: Model to use (default: gpt-4o)
            config: Optional full configuration
        """
        super().__init__(agent_id, model)
        self.api_key = api_key
        self.config = config or OpenAIAgentsConfig(api_key=api_key, model=model)
        self._client = None
        self._agent = None
        self._runner = None
        self._last_packet: Optional[StatusPacket] = None

    def _get_client(self):
        """Get or create the OpenAI client."""
        if self._client is None:
            try:
                # Try to import OpenAI Agents SDK
                from openai_agents import Agent, Runner

                self._agent = Agent(
                    name=self.agent_id,
                    instructions="You are a helpful coding assistant.",
                    model=self.config.model,
                )
                self._runner = Runner()
                self._client = "agents_sdk"

            except ImportError:
                # Fall back to direct OpenAI API
                try:
                    import openai

                    self._client = openai.OpenAI(api_key=self.api_key)
                except ImportError:
                    raise ImportError(
                        "Neither openai-agents nor openai package is installed. "
                        "Install with: pip install openai-agents or pip install openai"
                    )
        return self._client

    async def execute(self, task: str, context: dict[str, Any]) -> AgentResponse:
        """
        Execute a task using the OpenAI Agents SDK.

        Args:
            task: The task description/prompt
            context: Additional context including project state

        Returns:
            AgentResponse with result and usage metrics
        """
        self._status = AgentStatus.RUNNING
        start_time = datetime.now()

        try:
            client = self._get_client()

            # Build the prompt with context
            prompt = self._build_prompt(task, context)

            # Check if we're using the SDK or direct API
            if client == "agents_sdk":
                response = await self._execute_with_sdk(prompt)
            else:
                response = await self._execute_with_api(client, prompt)

            # Update usage stats
            self._usage_stats.add_response(response)

            # Create status packet
            self._last_packet = create_success_packet(
                agent_id=self.agent_id,
                task_id=self._current_task_id or "",
                summary=f"Completed: {task[:100]}...",
                files_modified=response.artifacts.files_modified,
                next_steps=["Review output", "Run tests"],
            )
            self._last_packet.artifacts = response.artifacts

            self._status = AgentStatus.IDLE
            return response

        except Exception as e:
            error_msg = SecretRedactor.redact(str(e))
            self._last_packet = create_failed_packet(
                agent_id=self.agent_id,
                task_id=self._current_task_id or "",
                error_message=error_msg,
            )
            self._status = AgentStatus.IDLE
            self._usage_stats.errors_count += 1

            return AgentResponse(
                content="",
                success=False,
                error=error_msg,
            )

    async def _execute_with_sdk(self, prompt: str) -> AgentResponse:
        """Execute using the OpenAI Agents SDK."""
        from openai_agents import Agent, Runner

        # Run the agent
        result = await self._runner.run(
            agent=self._agent,
            messages=[{"role": "user", "content": prompt}],
        )

        # Extract usage from result
        usage = getattr(result, "usage", None)
        if usage:
            tokens_input = getattr(usage, "prompt_tokens", 0)
            tokens_output = getattr(usage, "completion_tokens", 0)
            total_tokens = getattr(usage, "total_tokens", tokens_input + tokens_output)
        else:
            tokens_input = 0
            tokens_output = 0
            total_tokens = 0

        # Extract content
        content = ""
        if hasattr(result, "messages") and result.messages:
            content = result.messages[-1].get("content", "")
        elif hasattr(result, "content"):
            content = result.content

        return AgentResponse(
            content=content,
            tokens_input=tokens_input,
            tokens_output=tokens_output,
            tokens_used=total_tokens,
            cost=self._estimate_cost(tokens_input, tokens_output),
            model=self.model,
            artifacts=TaskArtifacts(
                tokens_input=tokens_input,
                tokens_output=tokens_output,
            ),
            success=True,
        )

    async def _execute_with_api(self, client, prompt: str) -> AgentResponse:
        """Execute using the direct OpenAI API."""
        # Run in thread pool since openai client is sync by default
        response = await asyncio.to_thread(
            client.chat.completions.create,
            model=self.config.model,
            max_tokens=self.config.max_tokens,
            messages=[{"role": "user", "content": prompt}],
        )

        # Extract content
        content = response.choices[0].message.content if response.choices else ""

        # Extract usage
        tokens_input = response.usage.prompt_tokens if response.usage else 0
        tokens_output = response.usage.completion_tokens if response.usage else 0

        return AgentResponse(
            content=content,
            tokens_input=tokens_input,
            tokens_output=tokens_output,
            tokens_used=tokens_input + tokens_output,
            cost=self._estimate_cost(tokens_input, tokens_output),
            model=response.model,
            artifacts=TaskArtifacts(
                tokens_input=tokens_input,
                tokens_output=tokens_output,
            ),
            success=True,
        )

    async def stream(self, task: str, context: dict[str, Any]) -> AsyncIterator[str]:
        """
        Stream response for real-time output.

        Args:
            task: The task description
            context: Additional context

        Yields:
            Response chunks as they become available
        """
        self._status = AgentStatus.RUNNING

        try:
            client = self._get_client()

            if client == "agents_sdk":
                # Fall back to non-streaming for SDK
                response = await self.execute(task, context)
                yield response.content
            else:
                # Use streaming API
                prompt = self._build_prompt(task, context)

                stream = client.chat.completions.create(
                    model=self.config.model,
                    max_tokens=self.config.max_tokens,
                    messages=[{"role": "user", "content": prompt}],
                    stream=True,
                )

                for chunk in stream:
                    if chunk.choices and chunk.choices[0].delta.content:
                        yield chunk.choices[0].delta.content

        except Exception as e:
            yield f"Error: {SecretRedactor.redact(str(e))}"

        finally:
            self._status = AgentStatus.IDLE

    def _build_prompt(self, task: str, context: dict[str, Any]) -> str:
        """Build the full prompt with context."""
        builder = PromptBuilder.for_api_agent()
        return builder.build(
            task,
            context,
            output_requirements=PromptBuilder.minimal_output_requirements(),
        )

    def _estimate_cost(self, tokens_input: int, tokens_output: int) -> float:
        """Estimate cost based on token usage."""
        return estimate_cost(self.model, tokens_input, tokens_output, provider="openai")

    def write_status_packet(self) -> StatusPacket:
        """Return the last status packet."""
        if self._last_packet:
            return self._last_packet

        return StatusPacket(
            agent_id=self.agent_id,
            task_id=self._current_task_id or "",
            status="idle",
            artifacts=TaskArtifacts(
                tokens_input=self._usage_stats.tokens_input,
                tokens_output=self._usage_stats.tokens_output,
                cost_usd=self._usage_stats.total_cost,
            ),
        )

    def is_healthy(self) -> bool:
        """Check if the adapter can connect to the API."""
        try:
            # Simple health check - verify API key format
            if not self.api_key or not self.api_key.startswith("sk-"):
                return False
            return True
        except Exception:
            return False

    def get_name(self) -> str:
        """Get adapter name."""
        return f"OpenAIAgents({self.model})"


class OpenAICodexAdapter(OpenAIAgentsAdapter):
    """
    Specialized adapter for Codex-style tasks.

    Uses GPT-4o by default but optimized for code generation tasks.
    """

    def __init__(
        self,
        agent_id: str,
        api_key: str,
        model: str = "gpt-4o",
    ) -> None:
        super().__init__(agent_id, api_key, model)

    def _build_prompt(self, task: str, context: dict[str, Any]) -> str:
        """Build prompt optimized for code generation."""
        base_prompt = super()._build_prompt(task, context)

        # Add code-specific instructions
        code_instructions = """
## Code Generation Guidelines
- Write clean, well-documented code
- Follow existing code style in the project
- Include error handling where appropriate
- Write tests for new functionality
- Use type hints in Python code
"""
        return base_prompt + "\n" + code_instructions

    def get_name(self) -> str:
        """Get adapter name."""
        return f"OpenAICodex({self.model})"
