from __future__ import annotations
from dataclasses import dataclass, field
from typing import Any, Callable

@dataclass
class Tool:
    name: str
    description: str
    parameters: dict[str, Any]
    fn: Callable[..., str]

    def to_schema(self) -> dict[str, Any]:
        return {
            "type": "function",
            "function": {
                "name": self.name,
                "description": self.description,
                "parameters": self.parameters,
            },
        }


class AgentCore:
    """
    Single-agent loop with a pluggable LLM client, tool registry,
    and swappable memory backend.
    """

    def __init__(
        self,
        llm_client: Any,
        model: str,
        system_prompt: str,
        tools: list[Tool] | None = None,
        memory_backend: Any | None = None,
        max_iterations: int = 25,
    ) -> None:
        self.llm = llm_client
        self.model = model
        self.system_prompt = system_prompt
        self.tools: dict[str, Tool] = {t.name: t for t in (tools or [])}
        self.memory = memory_backend
        self.max_iterations = max_iterations
        self._messages: list[dict[str, Any]] = [
            {"role": "system", "content": system_prompt}
        ]

    # ── Public API ──────────────────────────────────────────────────────────

    def run(self, goal: str) -> str:
        self._messages.append({"role": "user", "content": goal})
        self._inject_memory(goal)          # retrieve relevant past context

        for _ in range(self.max_iterations):
            response = self._call_llm()
            choice = response.choices[0]

            if choice.finish_reason == "tool_calls":
                for call in choice.message.tool_calls:
                    result = self._execute_tool(call)
                    self._messages.append({
                        "role": "tool",
                        "content": result,
                        "tool_call_id": call.id,
                    })

            elif choice.finish_reason == "stop":
                answer = choice.message.content
                self._persist_memory(goal, answer)   # store to long-term memory
                return answer

        raise RuntimeError("Max iterations reached without FINISH signal")

    # ── Private helpers ──────────────────────────────────────────────────────

    def _call_llm(self) -> Any:
        return self.llm.chat.completions.create(
            model=self.model,
            messages=self._messages,
            tools=[t.to_schema() for t in self.tools.values()] or None,
        )

    def _execute_tool(self, call: Any) -> str:
        tool = self.tools.get(call.function.name)
        if tool is None:
            return f"Error: unknown tool '{call.function.name}'"
        import json
        args = json.loads(call.function.arguments)
        return tool.fn(**args)

    def _inject_memory(self, query: str) -> None:
        if self.memory:
            relevant = self.memory.search(query, top_k=5)
            if relevant:
                self._messages.append({
                    "role": "system",
                    "content": "Relevant past context:\n" + "\n".join(relevant),
                })

    def _persist_memory(self, query: str, answer: str) -> None:
        if self.memory:
            self.memory.store(query=query, answer=answer)