from llm.bedrock import BedrockClient
from jarvis_mcp.client import MCPClient  # oder wo auch immer dein MCP-Client liegt
from config import MCP_URL
from util.logger import get_logger
import asyncio


class Agent:
    def __init__(self):
        self.logger = get_logger("NORA")
    
        self.llm = BedrockClient()
        self.mcp = MCPClient(MCP_URL)
    
        self.tools = []
    
        # System-Prompt als erste User-Nachricht (Nova 2 Lite kompatibel)
        self.messages = [
            {
                "role": "user",
                "content": [
                    {
                        "text": (
                            "Du bist N.O.R.A (Neural Operations for Residential Automation), ein freundlicher und kompetenter Smart-Home Assistent. Du bist ähnlich wie die KI Jarvis. höfflich, zielorientiert und leicht zynisch und ironisch. Du antwortest in kurzen Sätzen ausser es wird eine lange Erklärung verlangt. "
                            "für OpenHAB. Du steuerst Geräte und beantwortest Fragen über das Zuhause. "
                            "Nutze Tools, wenn sinnvoll. Antworte auf Deutsch und sei hilfreich."
                        )
                    }
                ]
            }
        ]
    # -------------------------
    # TOOL LOADING
    # -------------------------
    async def load_tools(self):
        try:
            res = await self.mcp.list_tools()
    
            self.tools = []
            for t in res.tools:
                schema = t.inputSchema
    
                # Bedrock-kompatibles Format erzwingen
                if isinstance(schema, dict) and "json" not in schema:
                    bedrock_schema = {"json": schema}
                else:
                    bedrock_schema = schema
    
                self.tools.append({
                    "toolSpec": {
                        "name": t.name,
                        "description": t.description,
                        "inputSchema": bedrock_schema
                    }
                })
    
            self.logger.info(f"{len(self.tools)} Tools erfolgreich für Bedrock angepasst und geladen")
    
        except Exception as e:
            self.logger.error(f"Tool Load Error: {e}")
            self.tools = []

    # -------------------------
    # TOOL EXECUTION
    # -------------------------
    async def _run_tool(self, tool):
        try:
            # Robustere Extraktion
            tool_use = tool.get("toolUse", tool)  # falls schon extrahiert
            name = tool_use.get("name")
            args = tool_use.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 ({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 iteration in range(10):  # etwas mehr Schleifendurchläufe erlaubt
            try:
                response = await self.llm.chat(
                    messages=self.messages,
                    tools=self.tools if self.tools else None
                )

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

            # Response-Struktur von Bedrock Converse
            output = response.get("output", {})
            message = output.get("message", {})
            content = message.get("content", [])

            tool_uses = []
            final_text = []

            # Verbessertes Parsing
            for c in content:
                if "toolUse" in c:
                    tool_uses.append(c)
                elif "text" in c:
                    final_text.append(c["text"])

            # TOOL EXECUTION PATH
            if tool_uses:
                self.logger.info(f"{len(tool_uses)} Tool(s) detected")

                # Assistant-Nachricht mit Tool-Call speichern
                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("toolUse", tool).get("toolUseId"),
                            "content": [{"text": str(result)}]
                        }
                    })

                # Tool-Ergebnisse zurück an das Modell
                self.messages.append({
                    "role": "user",
                    "content": tool_results
                })

                continue  # nächste Runde für finale Antwort

            # 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. Bitte versuche es erneut."