from __future__ import annotations

from dataclasses import dataclass
from typing import Any, Dict, List, Optional

from ..utils.logger import get_logger
from .orchestrator import Agent, AgentCapability, AgentType, Task


@dataclass
class LayerAgentConfig:
    """Config for a Layer-backed agent."""

    role: str = ""
    response_depth: str = "balanced"


class LayerAgent(Agent):
    """
    Agent implementation backed by a `ToxoLayer`.

    Supports task types:
    - `layer_query`: runs `await layer.query(question, context=..., response_depth=...)`
    """

    def __init__(
        self,
        *,
        agent_id: str,
        layer: Any,
        role: str = "",
        response_depth: str = "balanced",
        extra_capabilities: Optional[List[AgentCapability]] = None,
        config: Optional[Dict[str, Any]] = None,
    ):
        self.layer = layer
        self.layer_role = role
        self.default_response_depth = response_depth or "balanced"
        self.logger = get_logger(f"{__name__}.{agent_id}")

        capabilities = [
            AgentCapability(
                name="layer_query",
                description="Answer a question using a loaded ToxoLayer",
                input_types=["text"],
                output_types=["text"],
                max_concurrent_tasks=1,
                estimated_duration=15.0,
            )
        ]
        if extra_capabilities:
            capabilities.extend(extra_capabilities)

        super().__init__(
            agent_id=agent_id,
            agent_type=AgentType.SPECIALIST,
            capabilities=capabilities,
            config=config or {},
        )

    async def _setup(self):
        # Nothing heavy to initialize here; the underlying layer is already loaded.
        return None

    async def _cleanup(self):
        return None

    async def execute_task(self, task: Task) -> Dict[str, Any]:
        if task.type != "layer_query":
            raise ValueError(f"LayerAgent cannot execute task type: {task.type}")

        question = str(task.input_data.get("question") or "")
        if not question.strip():
            raise ValueError("layer_query task requires input_data.question")

        context = task.input_data.get("context")
        response_depth = task.input_data.get("response_depth") or self.default_response_depth

        # Prefix with agent role to reduce collisions when multiple agents answer.
        if self.layer_role:
            question = f"[Role: {self.layer_role}]\n\n{question}"

        response = await self.layer.query(question, context=context, response_depth=response_depth)
        return {"response": response}

