import json
import logging

from personal_agent.context import ContextBuilder
from personal_agent.identity import AgentIdentity, UserProfile
from personal_agent.llm import LLM
from personal_agent.memory.models import WorkingMemory
from personal_agent.memory.store import LongTermMemoryStore
from personal_agent.tools import (
    LongTermMemoryDeleteTool,
    LongTermMemorySaveTool,
    LongTermMemorySearchTool,
    MemoryAddFactTool,
    MemoryAddNoteTool,
    MemoryRemoveFactTool,
    MemoryRemoveNoteTool,
    TaskStartTool,
    ToolError,
    ToolRegistry,
)

logger = logging.getLogger(__name__)


class Agent:
    def __init__(
        self,
        llm: LLM,
        agent_identity: AgentIdentity,
        user_profile: UserProfile,
        tools: ToolRegistry,
        long_term_memory_store: LongTermMemoryStore,
    ) -> None:
        self.llm = llm
        self.tools = tools

        self.context = ContextBuilder(
            agent_identity,
            user_profile,
        )

        self.working_memory = WorkingMemory()
        self.long_term_memory_store = long_term_memory_store

        self.tools.register(MemoryAddFactTool(self.working_memory))
        self.tools.register(MemoryRemoveFactTool(self.working_memory))
        self.tools.register(MemoryRemoveNoteTool(self.working_memory))
        self.tools.register(MemoryAddNoteTool(self.working_memory))
        self.tools.register(TaskStartTool(self.working_memory))
        self.tools.register(LongTermMemorySaveTool(self.long_term_memory_store))
        self.tools.register(LongTermMemorySearchTool(self.long_term_memory_store))
        self.tools.register(LongTermMemoryDeleteTool(self.long_term_memory_store))

        self.messages: list[dict[str, object]] = []
        logger.info("Agent initialized")

    def chat(self, user_input: str) -> str:
        logger.info("User message received")
        self.messages.append(
            {
                "role": "user",
                "content": user_input,
            }
        )

        # print(f"[debug: {self.context.build_messages(self.messages, self.working_memory)}]")

        for _ in range(5):
            response = self.llm.chat(
                self.context.build_messages(
                    self.messages,
                    self.working_memory,
                ),
                tools=self.tools.schemas(),
            )

            if response.tool_calls:
                for call in response.tool_calls:
                    if call.name == "long_term_memory_save":
                        logger.info(
                            "Memory save requested: category=%s, importance=%s",
                            call.arguments.get("category"),
                            call.arguments.get("importance"),
                        )
                    else:
                        logger.info("Tool call requested: %s", call.name)

            if not response.tool_calls:
                answer = response.content or ""

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

                return answer

            self.messages.append(
                {
                    "role": "assistant",
                    "content": response.content,
                    "tool_calls": [
                        {
                            "id": call.id,
                            "type": "function",
                            "function": {
                                "name": call.name,
                                "arguments": json.dumps(call.arguments),
                            },
                        }
                        for call in response.tool_calls
                    ],
                }
            )

            for call in response.tool_calls:
                tool = self.tools.get(call.name)

                if tool is None:
                    result = f"Unknown tool: {call.name}"
                else:
                    try:
                        result = tool.execute(call.arguments).content
                        logger.info("Tool execution completed: %s", call.name)
                    except ToolError as error:
                        result = f"Tool error: {error}"
                        logger.info("Tool execution failed: %s", call.name)

                self.messages.append(
                    {
                        "role": "tool",
                        "tool_call_id": call.id,
                        "content": result,
                    }
                )
        raise RuntimeError("Tool call limit exceeded")

    def end_task(self) -> None:
        self.working_memory.end_task()
        logger.info("Task ended")

    def start_task(self, objective: str) -> None:
        self.working_memory.start_task(objective)
        logger.info("Task started")

    def pause_task(self) -> None:
        self.working_memory.pause_task()
        logger.info("Task paused")

    def resume_task(self) -> None:
        self.working_memory.resume_task()
        logger.info("Task resumed")

    def complete_task(self) -> None:
        self.working_memory.complete_task()
        logger.info("Taks completed")