"""
LangChain-based base agent.

Provides a simple interface around LangChain's initialize_agent, with
pluggable LLM backends and StructuredTool support.
"""

from __future__ import annotations

from typing import List, Optional, Any
import os

from langchain.agents import initialize_agent, AgentType
from langchain.tools import BaseTool


def _default_model():
    """Return a default LangChain chat model based on env LLM_BACKEND."""
    backend = (os.environ.get("LLM_BACKEND") or "openai").lower()
    if backend == "gemini":
        try:
            from langchain_google_genai import ChatGoogleGenerativeAI
            api_key = os.environ.get("GEMINI_API_KEY")
            return ChatGoogleGenerativeAI(model="gemini-2.5-flash", google_api_key=api_key, temperature=0)
        except Exception as e:
            raise RuntimeError(f"Gemini backend requested but not available: {e}")
    elif backend == "openai":
        try:
            from langchain_openai import ChatOpenAI
            return ChatOpenAI(model="gpt-4o-mini", temperature=0)
        except Exception as e:
            raise RuntimeError(f"OpenAI backend requested but not available: {e}")
    else:
        # Fallback to a minimal echo-like model using a simple wrapper
        class _Echo:
            async def ainvoke(self, inputs: Any):
                msg = inputs.get("input") if isinstance(inputs, dict) else str(inputs)
                return {"output": msg}

        return _Echo()


class BaseAgent:
    def __init__(
        self,
        name: str,
        description: str = "",
        tools: Optional[List[BaseTool]] = None,
        model: Any = None,
        allow_delegation: bool = True,
    ):
        self.name = name
        self.description = description
        self.allow_delegation = allow_delegation
        self.tools: List[BaseTool] = tools or []
        self.model = model or _default_model()
        self._init_agent()

    def _init_agent(self) -> None:
        # initialize_agent expects a list of tools and a model implementing LC API
        self.agent = initialize_agent(
            self.tools,
            self.model,
            agent_type=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
            verbose=False,
            handle_parsing_errors=True,
        )

    async def handle(self, task: dict):
        """Handle a single task using LangChain."""
        query = task.get("query") or task.get("message") or str(task)
        return await self.agent.ainvoke({"input": query})

    def bind_tools(self, new_tools: List[BaseTool]):
        """Bind additional tools dynamically and reinitialize agent."""
        self.tools.extend(new_tools)
        self._init_agent()
