from llm.bedrock import BedrockClient
from jarvis_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. "
                            "Du steuerst Geräte über Tools. "
                            "Nutze Tools wenn sinnvoll."
                        )
                    }
                ]
            }
        ]

    # -------------------------
    # TOOL LOADING
    # -------------------------
    async def load_tools(self):
        try:
            res = await self.mcp.list_tools()

            self.tools = [
                {
                    "toolSpec": {
                        "name": t.name,
                        "description": t.description,
                        "inputSchema": t.inputSchema
                    }
                }
                for t in res.tools
            ]

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

        except Exception as e:
            self.logger.error(f"Tool Load Error: {e}")
            self.tools = []

    # -------------------------
    # TOOL EXECUTION
    # -------------------------
    async def _run_tool(self, tool):
        try:
            name = tool.get("name")
            args = tool.get("input", {}) or {}

            self.logger.info(f"Tool Call → {name} | args={args}")

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

            # MCP result stabilisieren
            if hasattr(result, "content"):
                return result.content

            return result

        except Exception as e:
            self.logger.error(f"Tool Error ({tool.get('name')}): {e}")
            return f"ERROR: {e}"

    # -------------------------
    # MAIN LOOP
    # -------------------------
    async def run(self, user_input: str):

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

        for _ in range(8):

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

            except Exception as e:
                self.logger.error(f"LLM Error: {e}")
                return f"LLM Fehler: {e}"

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

            tool_uses = []
            final_text = []

            # -------------------------
            # PARSE RESPONSE SAFE
            # -------------------------
            for c in content:

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

                if "text" in c:
                    final_text.append(c["text"])

            # -------------------------
            # TOOL EXECUTION PATH
            # -------------------------
            if tool_uses:

                self.logger.info(f"{len(tool_uses)} Tool(s) detected")

                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.get("toolUseId"),
                            "content": [
                                {
                                    "text": str(result)
                                }
                            ]
                        }
                    })

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

                continue

            # -------------------------
            # FINAL RESPONSE
            # -------------------------
            if final_text:

                text = "\n".join(final_text)

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

                return text

        return "Tool Loop Limit erreicht"