from llm.bedrock import BedrockClient
from mcp.client import MCPClient
from config import MCP_URL
from util.logger import get_logger


class Agent:
    def __init__(self):
        self.logger = get_logger("NORA")

        self.llm = BedrockClient()
        self.mcp = MCPClient(MCP_URL)

        self.tools = []
        self.messages = [
            {
                "role": "system",
                "content": [
                    {
                        "text": "Du bist N.O.R.A, ein Smart-Home Assistent. Nutze Tools wenn nötig."
                    }
                ]
            }
        ]

    async def load_tools(self):
        res = await self.mcp.list_tools()

        self.tools = []

        for t in res.tools:
            self.tools.append({
                "toolSpec": {
                    "name": t.name,
                    "description": t.description,
                    "inputSchema": t.inputSchema
                }
            })

        self.logger.info(f"{len(self.tools)} Tools geladen")

    async def _run_tool(self, tool):
        name = tool["name"]
        args = tool.get("input", {})

        self.logger.info(f"Tool Call: {name} -> {args}")

        result = await self.mcp.call_tool(name, args)

        return result

    async def run(self, user_input: str):

        self.messages.append({
            "role": "user",
            "content": [{"text": user_input}]
        })

        for _ in range(8):

            response = await self.llm.chat(
                messages=self.messages,
                tools=self.tools
            )

            output = response.get("output", {})
            message = output.get("message", {})
            content = message.get("content", [])

            tool_uses = []
            final_text = ""

            for c in content:
                tool = c.get("toolUse")
                if tool:
                    tool_uses.append(tool)

                if "text" in c:
                    final_text += c["text"]

            if tool_uses:

                self.messages.append({
                    "role": "assistant",
                    "content": content
                })

                tool_results = []

                for tool in tool_uses:
                    result = await self._run_tool(tool)

                    tool_results.append({
                        "toolResult": {
                            "toolUseId": tool["toolUseId"],
                            "content": [
                                {
                                    "text": str(result)
                                }
                            ]
                        }
                    })

                self.messages.append({
                    "role": "user",
                    "content": tool_results
                })

                continue

            if final_text:
                self.messages.append({
                    "role": "assistant",
                    "content": [{"text": final_text}]
                })

                return final_text

        return "Tool Loop Limit erreicht"