"""
llm_agents.base_agent
======================
BaseAgent: a thin wrapper around litellm.completion that manages:
  - conversation history
  - tool/function calling (OpenAI-style JSON schema)
  - streaming responses
  - automatic retries with exponential back-off

Usage
-----
    from ml_ai_library.llm_agents.base_agent import BaseAgent, Tool

    def add(a: int, b: int) -> int:
        return a + b

    agent = BaseAgent(
        model="gpt-4o-mini",          # any litellm model string
        system="You are a calculator.",
        tools=[Tool.from_fn(add)],
    )
    reply = agent.chat("What is 42 + 58?")
    print(reply)   # "The sum is 100."
"""

from __future__ import annotations
import inspect
import json
import time
from dataclasses import dataclass, field
from typing import Any, Callable, Dict, List, Optional

import litellm


# ---------------------------------------------------------------------------
# Tool definition
# ---------------------------------------------------------------------------

@dataclass
class Tool:
    """Wraps a Python callable as an OpenAI-style tool schema."""

    name: str
    description: str
    fn: Callable
    parameters: Dict

    @classmethod
    def from_fn(cls, fn: Callable, description: str = "") -> "Tool":
        """
        Auto-build a Tool from a Python function's signature and docstring.
        Only supports str / int / float / bool parameter types.
        """
        sig  = inspect.signature(fn)
        desc = description or (fn.__doc__ or "").strip().split("\n")[0]
        props: Dict[str, Any] = {}
        required: List[str]   = []
        type_map = {str: "string", int: "integer", float: "number", bool: "boolean"}

        for pname, param in sig.parameters.items():
            ann   = param.annotation
            jtype = type_map.get(ann, "string")
            props[pname] = {"type": jtype}
            if param.default is inspect.Parameter.empty:
                required.append(pname)

        return cls(
            name=fn.__name__,
            description=desc,
            fn=fn,
            parameters={"type": "object", "properties": props,
                        "required": required},
        )

    def schema(self) -> Dict:
        return {
            "type": "function",
            "function": {
                "name": self.name,
                "description": self.description,
                "parameters": self.parameters,
            },
        }

    def call(self, **kwargs) -> Any:
        return self.fn(**kwargs)


# ---------------------------------------------------------------------------
# BaseAgent
# ---------------------------------------------------------------------------

class BaseAgent:
    """
    Conversational LLM agent backed by litellm.

    Parameters
    ----------
    model      : litellm model string, e.g. "gpt-4o", "claude-3-5-sonnet",
                 "ollama/llama3", "gemini/gemini-pro".
    system     : system prompt.
    tools      : list of Tool objects available to the model.
    max_turns  : maximum tool-call iterations per user message.
    temperature: sampling temperature.
    """

    def __init__(
        self,
        model: str = "gpt-4o-mini",
        system: str = "You are a helpful assistant.",
        tools: Optional[List[Tool]] = None,
        max_turns: int = 10,
        temperature: float = 0.0,
    ) -> None:
        self.model       = model
        self.tools       = {t.name: t for t in (tools or [])}
        self.max_turns   = max_turns
        self.temperature = temperature
        self.history: List[Dict] = []
        if system:
            self.history.append({"role": "system", "content": system})

    # ------------------------------------------------------------------
    # Public interface
    # ------------------------------------------------------------------

    def chat(self, user_message: str) -> str:
        """Send a message and return the final text reply."""
        self.history.append({"role": "user", "content": user_message})
        return self._run_loop()

    def reset(self, keep_system: bool = True):
        """Clear conversation history."""
        sys_msgs = [m for m in self.history if m["role"] == "system"]
        self.history = sys_msgs if keep_system else []

    # ------------------------------------------------------------------
    # Internal loop
    # ------------------------------------------------------------------

    def _run_loop(self) -> str:
        tool_schemas = [t.schema() for t in self.tools.values()] or None
        for _ in range(self.max_turns):
            response = self._call_llm(tool_schemas)
            msg      = response.choices[0].message

            # Plain text reply
            if not getattr(msg, "tool_calls", None):
                text = msg.content or ""
                self.history.append({"role": "assistant", "content": text})
                return text

            # Tool call(s)
            self.history.append(msg.model_dump())
            for tc in msg.tool_calls:
                fn_name = tc.function.name
                fn_args = json.loads(tc.function.arguments)
                tool    = self.tools.get(fn_name)
                if tool is None:
                    result = f"[Error: tool '{fn_name}' not found]"
                else:
                    try:
                        result = str(tool.call(**fn_args))
                    except Exception as e:
                        result = f"[Error: {e}]"
                self.history.append({
                    "role": "tool",
                    "tool_call_id": tc.id,
                    "content": result,
                })

        return "[Max turns reached]"

    def _call_llm(self, tools=None, retries: int = 3) -> Any:
        kwargs: Dict[str, Any] = {
            "model":       self.model,
            "messages":    self.history,
            "temperature": self.temperature,
        }
        if tools:
            kwargs["tools"] = tools

        delay = 2.0
        for attempt in range(retries):
            try:
                return litellm.completion(**kwargs)
            except Exception as e:
                if attempt == retries - 1:
                    raise
                time.sleep(delay)
                delay *= 2
