"""ORQ backend implementation for dynamic red teaming.

Consolidates all ORQ SDK-specific agent code behind the backend protocols
defined in ``backends.base``. Other modules import these concrete classes
when they need ORQ-specific behavior, or accept the protocols when they
want to be backend-agnostic.
"""

from __future__ import annotations

import asyncio
from typing import TYPE_CHECKING, Any

import httpx
from loguru import logger

try:
    from orq_ai_sdk import Orq
    from orq_shared.config import get_config  # pyright: ignore[reportMissingImports]
except ImportError:
    Orq = None  # type: ignore[assignment,misc]
    get_config = None  # type: ignore[assignment]

from evaluatorq.redteam.backends.base import extract_provider_error_code, extract_status_code
from evaluatorq.redteam.contracts import (
    PIPELINE_CONFIG,
    AgentContext,
    KnowledgeBaseInfo,
    MemoryStoreInfo,
    TokenUsage,
    ToolInfo,
)

if TYPE_CHECKING:
    from orq_ai_sdk import Orq as OrqClient

    from evaluatorq.redteam.backends.base import AgentTarget


class ORQAgentTarget:
    """Target adapter for ORQ agents.

    Wraps the ORQ SDK to provide the AgentTarget protocol.
    """

    def __init__(
        self,
        agent_key: str,
        orq_client: OrqClient,
        memory_entity_id: str | None = None,
    ):
        self.agent_key: str = agent_key
        self.orq_client: OrqClient = orq_client
        self.memory_entity_id: str | None = memory_entity_id
        self._task_id: str | None = None
        self._last_token_usage: TokenUsage | None = None

    async def send_prompt(self, prompt: str) -> str:
        """Send a prompt to the ORQ agent."""
        try:
            kwargs: dict[str, Any] = {
                'agent_key': self.agent_key,
                'message': {'role': 'user', 'parts': [{'kind': 'text', 'text': prompt}]},
                'task_id': self._task_id,
                'background': False,
            }
            if self.memory_entity_id:
                kwargs['memory'] = {'entity_id': self.memory_entity_id}

            response = await asyncio.to_thread(self.orq_client.agents.responses.create, **kwargs)

            resp_task_id = getattr(response, 'task_id', None)
            if resp_task_id:
                self._task_id = resp_task_id

            total_prompt_tokens = 0
            total_completion_tokens = 0
            total_tokens = 0
            total_calls = 0

            def _accumulate_usage(resp: object) -> None:
                nonlocal total_prompt_tokens, total_completion_tokens, total_tokens, total_calls
                usage = getattr(resp, 'usage', None)
                if usage is None:
                    return
                prompt_tokens = int(getattr(usage, 'prompt_tokens', 0) or 0)
                completion_tokens = int(getattr(usage, 'completion_tokens', 0) or 0)
                total = int(getattr(usage, 'total_tokens', prompt_tokens + completion_tokens) or 0)
                total_prompt_tokens += prompt_tokens
                total_completion_tokens += completion_tokens
                total_tokens += total
                total_calls += 1

            def _extract_text(resp: object) -> str:
                output = getattr(resp, 'output', None) or []
                for item in output:
                    parts = getattr(item, 'parts', None) or []
                    for part in parts:
                        if getattr(part, 'kind', None) == 'text':
                            text = getattr(part, 'text', None)
                            if isinstance(text, str) and text.strip():
                                return text
                return ''

            def _pending_tool_call_ids(resp: object) -> list[str]:
                pending = getattr(resp, 'pending_tool_calls', None) or []
                ids: list[str] = []
                for call in pending:
                    call_id = getattr(call, 'id', None)
                    if not call_id and isinstance(call, dict):
                        call_id = call.get('id')
                    if isinstance(call_id, str) and call_id.strip():
                        ids.append(call_id)
                return ids

            _accumulate_usage(response)
            text_response = _extract_text(response)

            # Some agent tool flows require client-provided tool_result parts.
            # Continue the same task with synthetic tool results so the thread can progress.
            max_tool_continuations = 5
            pending_ids = _pending_tool_call_ids(response)
            continuation_count = 0
            while pending_ids and continuation_count < max_tool_continuations:
                continuation_count += 1
                logger.debug(
                    f'{self.agent_key}: resolving {len(pending_ids)} pending tool call(s) via synthetic tool_result (step {continuation_count}/{max_tool_continuations})'
                )

                tool_parts = [
                    {
                        'kind': 'tool_result',
                        'tool_call_id': tool_call_id,
                        'result': {
                            'ok': False,
                            'error': 'Tool execution unavailable in red-teaming harness',
                        },
                    }
                    for tool_call_id in pending_ids
                ]
                response = await asyncio.to_thread(
                    self.orq_client.agents.responses.create,
                    agent_key=self.agent_key,
                    message={'role': 'tool', 'parts': tool_parts},
                    task_id=self._task_id,
                    background=False,
                )
                resp_task_id = getattr(response, 'task_id', None)
                if resp_task_id:
                    self._task_id = resp_task_id
                _accumulate_usage(response)
                extracted = _extract_text(response)
                if extracted:
                    text_response = extracted
                pending_ids = _pending_tool_call_ids(response)

            if pending_ids:
                raise RuntimeError(
                    f'Unresolved pending tool calls after {max_tool_continuations} continuations: {pending_ids}'
                )

            if total_calls > 0:
                self._last_token_usage = TokenUsage(
                    prompt_tokens=total_prompt_tokens,
                    completion_tokens=total_completion_tokens,
                    total_tokens=total_tokens,
                    calls=total_calls,
                )
            else:
                self._last_token_usage = None

            if text_response:
                return text_response

            return ''

        except Exception as e:
            logger.error(f'ORQ agent call failed: {e}')
            raise

    def reset_conversation(self) -> None:
        """Reset conversation state for a new attack."""
        self._task_id = None
        self._last_token_usage = None

    def consume_last_token_usage(self) -> TokenUsage | None:
        """Return and clear usage from the last send_prompt() call."""
        usage = self._last_token_usage
        self._last_token_usage = None
        return usage


class ORQContextProvider:
    """Retrieves agent context from the ORQ API."""

    def __init__(self, orq_client: OrqClient):
        self.orq_client: OrqClient = orq_client

    async def get_agent_context(self, agent_key: str) -> AgentContext:
        """Retrieve full agent context from ORQ API."""
        logger.info(f'Retrieving agent context for: {agent_key}')

        agent_data = await asyncio.to_thread(
            self.orq_client.agents.retrieve,
            agent_key=agent_key,
        )

        # Parse tools from settings.tools
        tools: list[ToolInfo] = []
        settings = getattr(agent_data, 'settings', None)
        if settings and hasattr(settings, 'tools') and settings.tools:
            tools.extend(
                ToolInfo(
                    name=getattr(tool, 'key', None) or getattr(tool, 'display_name', None) or tool.id,
                    description=getattr(tool, 'description', None),
                    parameters=None,
                )
                for tool in settings.tools
            )

        # Extract raw IDs for enrichment
        raw_kb_ids: list[str] = []
        if hasattr(agent_data, 'knowledge_bases') and agent_data.knowledge_bases:
            raw_kb_ids = [getattr(kb, 'knowledge_id', None) or str(kb) for kb in agent_data.knowledge_bases]

        raw_ms_ids: list[str] = []
        if hasattr(agent_data, 'memory_stores') and agent_data.memory_stores:
            raw_ms_ids = [ms if isinstance(ms, str) else getattr(ms, 'key', str(ms)) for ms in agent_data.memory_stores]  # pyright: ignore[reportUnnecessaryIsInstance]

        # Enrich knowledge bases and memory stores concurrently
        security = getattr(self.orq_client, 'sdk_configuration', None)
        security = getattr(security, 'security', None) if security else None
        api_key: str = getattr(security, 'api_key', '') if security else ''
        config = get_config() if get_config is not None else None
        server_url = config.orq_server_url if config else 'https://api.orq.ai'

        enrichment_tasks: list[Any] = [self._enrich_knowledge_base(kb_id) for kb_id in raw_kb_ids]
        enrichment_tasks.extend(self._enrich_memory_store(api_key, server_url, ms_id) for ms_id in raw_ms_ids)

        enriched_results: list[Any] = list(await asyncio.gather(*enrichment_tasks)) if enrichment_tasks else []

        knowledge_bases: list[KnowledgeBaseInfo] = [r for r in enriched_results if isinstance(r, KnowledgeBaseInfo)]
        memory_stores: list[MemoryStoreInfo] = [r for r in enriched_results if isinstance(r, MemoryStoreInfo)]

        model_raw = getattr(agent_data, 'model', None)
        model_id = getattr(model_raw, 'id', None) if model_raw is not None else None

        context = AgentContext(
            key=agent_key,
            display_name=getattr(agent_data, 'display_name', None),
            description=getattr(agent_data, 'description', None),
            system_prompt=getattr(agent_data, 'system_prompt', None),
            instructions=getattr(agent_data, 'instructions', None),
            tools=tools,
            memory_stores=memory_stores,
            knowledge_bases=knowledge_bases,
            model=model_id,
        )

        logger.info(
            f'Retrieved context: {len(tools)} tools, {len(memory_stores)} memory stores, {len(knowledge_bases)} knowledge bases'
        )

        return context

    async def _enrich_knowledge_base(self, kb_id: str) -> KnowledgeBaseInfo:
        """Retrieve full knowledge base details from ORQ API."""
        try:
            kb = await asyncio.to_thread(self.orq_client.knowledge.retrieve, knowledge_id=kb_id)
            return KnowledgeBaseInfo(
                id=kb_id,
                key=getattr(kb, 'key', None),
                name=getattr(kb, 'key', None),
                description=getattr(kb, 'description', None) or None,
            )
        except Exception as e:
            logger.warning(f'Failed to enrich knowledge base {kb_id}: {e}')
            return KnowledgeBaseInfo(id=kb_id)

    async def _enrich_memory_store(self, api_key: str, server_url: str, ms_id: str) -> MemoryStoreInfo:
        """Retrieve full memory store details from ORQ API."""
        try:
            async with httpx.AsyncClient() as client:
                r = await client.get(
                    f'{server_url}/v2/memory-stores/{ms_id}',
                    headers={'Authorization': f'Bearer {api_key}'},
                    timeout=10,
                )
                r.raise_for_status()
                data = r.json()
                return MemoryStoreInfo(
                    id=ms_id,
                    key=data.get('key'),
                    description=data.get('description') or None,
                )
        except Exception as e:
            logger.warning(f'Failed to enrich memory store {ms_id}: {e}')
            return MemoryStoreInfo(id=ms_id)


def _require_orq_sdk() -> None:
    """Raise a clear error if the ORQ SDK is not installed."""
    if Orq is None or get_config is None:
        raise RuntimeError(
            "ORQ SDK not installed. Install with: pip install 'evaluatorq[orq]'"
        )


class ORQTargetFactory:
    """Creates ORQAgentTarget instances, one per job."""

    def __init__(self, orq_client: OrqClient | None = None):
        self._orq_client: OrqClient
        if orq_client is not None:
            self._orq_client = orq_client
        else:
            _require_orq_sdk()
            # After _require_orq_sdk(), Orq and get_config are guaranteed non-None.
            config = get_config()  # pyright: ignore[reportOptionalCall]
            self._orq_client = Orq(  # pyright: ignore[reportOptionalCall]
                api_key=config.orq_api_key,
                server_url=config.orq_server_url,
                timeout_ms=PIPELINE_CONFIG.target_agent_timeout_ms,
            )

    def create_target(self, agent_key: str, memory_entity_id: str | None = None) -> AgentTarget:
        """Create a new ORQAgentTarget for the given agent key."""
        return ORQAgentTarget(
            agent_key=agent_key,
            orq_client=self._orq_client,
            memory_entity_id=memory_entity_id,
        )


class ORQMemoryCleanup:
    """Cleans up memory entities created during red teaming via ORQ API."""

    def __init__(self) -> None:
        _require_orq_sdk()
        # After _require_orq_sdk(), get_config is guaranteed non-None.
        config = get_config()  # pyright: ignore[reportOptionalCall]
        self._api_key: str = config.orq_api_key
        self._server_url: str = config.orq_server_url

    async def cleanup_memory(self, agent_context: AgentContext, entity_ids: list[str]) -> None:
        """Delete memory entities for each memory store x entity_id combination."""
        headers = {'Authorization': f'Bearer {self._api_key}'}
        failed = 0
        total = 0

        async with httpx.AsyncClient(timeout=10) as client:
            for ms in agent_context.memory_stores:
                if not ms.key:
                    logger.warning(f'Memory store {ms.id} has no key, skipping cleanup')
                    continue
                for entity_id in entity_ids:
                    total += 1
                    url = f'{self._server_url}/v2/memory-stores/{ms.key}/memories/{entity_id}'
                    try:
                        r = await client.delete(url, headers=headers)
                        if r.status_code == 204:
                            logger.debug(f'Deleted memory entity {entity_id} from store {ms.key}')
                        elif r.status_code == 404:
                            logger.debug(f'Memory entity {entity_id} not found in store {ms.key}')
                        else:
                            failed += 1
                            logger.warning(f'Memory cleanup for {ms.key}/{entity_id} returned {r.status_code}')
                    except Exception as e:
                        failed += 1
                        logger.warning(f'Failed to cleanup memory entity {entity_id} from {ms.key}: {e}')

        if failed == total and total > 0:
            logger.error(f'Memory cleanup failed: all {failed} operations failed')
        elif failed > 0:
            logger.warning(
                f'Memory cleanup partially failed: {failed}/{total} operations failed ({len(entity_ids)} entities across {len(agent_context.memory_stores)} stores)'
            )
        else:
            logger.info(
                f'Memory cleanup complete ({len(entity_ids)} entities across {len(agent_context.memory_stores)} stores)'
            )


class ORQErrorMapper:
    """Normalize ORQ SDK/HTTP failures into runtime error taxonomy."""

    def map_error(self, exc: Exception) -> tuple[str, str]:
        name = type(exc).__name__.lower()
        text = str(exc).lower()
        status_code = extract_status_code(exc)
        provider_code = extract_provider_error_code(exc)

        if status_code is not None:
            return f'orq.http.{status_code}', f'{type(exc).__name__}: {exc}'
        if provider_code:
            return f'orq.code.{provider_code}', f'{type(exc).__name__}: {exc}'
        if 'timeout' in name or 'timed out' in text:
            return 'orq.timeout', f'{type(exc).__name__}: {exc}'
        if 'auth' in name or 'unauthorized' in text or 'forbidden' in text:
            return 'orq.auth', f'{type(exc).__name__}: {exc}'
        if 'ratelimit' in name or '429' in text:
            return 'orq.rate_limit', f'{type(exc).__name__}: {exc}'
        return 'orq.unknown', f'{type(exc).__name__}: {exc}'


def create_orq_backend(
    orq_client: OrqClient | None = None,
) -> tuple[ORQTargetFactory, ORQContextProvider, ORQMemoryCleanup]:
    """Convenience function returning all three ORQ backend components.

    Args:
        orq_client: Optional pre-configured ORQ SDK client. If None, one is created
                    from environment config.

    Returns:
        Tuple of (target_factory, context_provider, memory_cleanup)
    """
    if orq_client is None:
        _require_orq_sdk()
        # After _require_orq_sdk(), Orq and get_config are guaranteed non-None.
        config = get_config()  # pyright: ignore[reportOptionalCall]
        orq_client = Orq(  # pyright: ignore[reportOptionalCall]
            api_key=config.orq_api_key,
            server_url=config.orq_server_url,
            timeout_ms=PIPELINE_CONFIG.target_agent_timeout_ms,
        )

    return (
        ORQTargetFactory(orq_client),
        ORQContextProvider(orq_client),
        ORQMemoryCleanup(),
    )
