from abc import ABC, abstractmethod
import json, time
import logging

logger = logging.getLogger(__name__)

class BaseLLM(ABC):
    def clear_history(self) -> None:
        self.message_history = []

    @abstractmethod
    async def get_reply(self, prompt: str):
        pass

    def append_user_message(self, text: str):
        self.message_history.append({"role": "user", "content": text})

    def append_assistant_message(self, text: str):
        self.message_history.append({"role": "assistant", "content": text})

class BaseAgent(ABC):
    def clear_history(self) -> None:
        self.message_history = []

    @abstractmethod
    async def get_reply_frontend(self, prompt: str):
        pass

    def append_user_message(self, text: str):
        self.message_history.append({"role": "user", "content": text})

    def append_assistant_message(self, text: str):
        self.message_history.append({"role": "assistant", "content": text})

    async def process_tool_use(self, tool_input, tool_name, user_chat_id):
        try:
            if tool_name == "retriever":
                tool_input["user_chat_id"] = user_chat_id

            result = await self.tools[tool_name][1](**tool_input)

            if tool_name == "retriever":
                scores = [res.score for res in result[1]]
                payloads = [res.payload for res in result[1]]
                output = "\n".join([
                    f"<doc file=\"{_get_key(payload)}\">{result[0][i]}"
                    for i, payload in enumerate(payloads)
                ])
                metadata = json.dumps({
                    "documents": [{
                        "key": _get_key(payload),
                        "filetype": payload.get("filetype", ""),
                        "score": scores[i],
                    } for i, payload in enumerate(payloads)]
                }, ensure_ascii=False)
                return output, metadata
            return result, None
        except Exception as e:
            import traceback
            logger.warning("ERROR tool result for %s: %s", tool_name, e)
            traceback.print_exc()
            return "ERROR", None


def _get_key(payload: dict) -> str:
    """Extract display key from payload, removing uploads/{userID} prefix."""
    key = payload.get("key")
    if key:
        return "/".join(key.split("/")[2:])
    file_dir = payload.get("fileDirectory", "")
    return "/".join(file_dir.split("/")[2:])