import os
from pathlib import Path
from typing import Dict, Any, List
import json
from litellm import completion
from abc import ABC, abstractmethod
import traceback
from src import constants

from src.tools.map import TOOL_MAP
from src import utils, constants
import tiktoken

ENC = tiktoken.get_encoding("cl100k_base")


class ToolOutputError(Exception):
    """Custom exception raised when a tool returns a known failure string."""

    pass


class BaseAgent(ABC):
    def __init__(
        self,
        model_name: str,
        temperature: float,
        sandbox_dir: str,
        pdb_id: str | None = None,
        ligand_name: str | None = None,
        model_supports_system_messages: bool = True,
    ):
        self.model_name = model_name
        self.temperature = temperature
        self.sandbox_dir = Path(sandbox_dir)
        self.pdb_id = pdb_id
        self.ligand_name = ligand_name

        self.model_supports_system_messages = model_supports_system_messages

        self.tool_schemas = None
        self.messages: List[Dict[str, Any]] = []
        self.llm_cost = 0

        self.logger = utils.get_class_logger(self.__class__.__name__)

    @abstractmethod
    def setup_tools(self):
        pass

    @abstractmethod
    def _additional_check_for_errors_tool_output(self, tool_name, tool_call) -> bool:
        pass

    import json

    def _count_tokens(self, text):
        if not isinstance(text, str):
            text = json.dumps(text)
        return len(ENC.encode(text))
    
    def _find_recent_block(self, messages):
        """
        Return the smallest suffix of messages that forms a logically valid block.
        """
        if not messages:
            return []

        # Start from the end
        block = []
        i = len(messages) - 1

        while i >= 0:
            m = messages[i]
            block.insert(0, m)

            role = m.get("role")

            # Stop conditions for different last-message types:

            # 1. If the last message is a 'tool' message, we must also include the
            #    assistant tool call that triggered it.
            if role == "tool":
                # include until we find the matching assistant tool call
                i -= 1
                while i >= 0 and messages[i].get("role") != "assistant":
                    block.insert(0, messages[i])
                    i -= 1
                if i >= 0:
                    block.insert(0, messages[i])
                return block

            # 2. If last message is assistant tool call (function call),
            #    we keep only that one so the model knows what it's continuing.
            if role == "assistant" and m.get("tool_calls"):
                return block

            # 3. If the last message is a normal assistant message,
            #    keep it and the preceding user message.
            if role == "assistant":
                # include the previous user message
                if i - 1 >= 0 and messages[i-1].get("role") == "user":
                    block.insert(0, messages[i-1])
                return block

            # 4. If the last message is a user message,
            #    include previous assistant message.
            if role == "user":
                if i - 1 >= 0 and messages[i-1].get("role") == "assistant":
                    block.insert(0, messages[i-1])
                return block

            i -= 1

        return block


    def _summarize_history(self, messages):
        recent_messages = self._find_recent_block(messages)
        history_to_summarize = messages[:-len(recent_messages)]
        
        # Turn history into text
        history_text = "\n".join(
            f"{m['role']}: {m['content']}"
            for m in history_to_summarize
        )

        # Summarize
        summary_response = completion(
            model=self.model_name,
            temperature=0.1,
            messages=[
                {"role": "system", "content": "Summarize the conversation concisely but fully."},
                {"role": "user", "content": history_text}
            ],
            max_tokens=constants.SUMMARY_OUTPUT_TOKENS,
        )

        summary_text = summary_response.choices[0].message["content"]

        # NEW conversation = summary + recent logical block
        return [
            {"role": "assistant", "content": f"[Conversation Summary]\n{summary_text}"},
            *recent_messages
        ]


    def _validate_tool_path(self, tool_input) -> None:
        # if "path" in tool_input and not utils.is_path_child_dir(tool_input["path"], self.sandbox_dir):
        #     raise PermissionError(f"Access outside sandbox not allowed: {tool_input['path']}, {self.sandbox_dir}")

        # if "sandbox_dir" in tool_input and not utils.is_path_child_dir(tool_input["sandbox_dir"], self.sandbox_dir):
            # raise PermissionError(f"Access outside sandbox not allowed: {tool_input['sandbox_dir']}, {self.sandbox_dir}")
        pass

    def _safe_execute_tool(self, tool_name: str, tool_input: Dict[str, Any]) -> dict:
        """Executes a tool and catches any errors to pass back to the LLM."""

        self.logger.info(f"Executing tool: {tool_name} with input: {tool_input}")

        if not isinstance(tool_input, dict):
            try:
                tool_input = json.loads(tool_input)
            except Exception:
                error_msg = f"Invalid tool input format (not a dict): {tool_input}"
                self.logger.error(error_msg)
                return {"ok": False, "output": error_msg}

        tool_output = None

        # try:
        self._validate_tool_path(tool_input)
        func = TOOL_MAP.get(tool_name)

        if not func:
            raise ValueError(f"Unknown tool: {tool_name}")

        tool_output = func(self, tool_input)
        passed = self._additional_check_for_errors_tool_output(tool_name, tool_output)

        return {"ok": passed, "output": tool_output}
        # except Exception as e:
        #     tb = traceback.format_exc()
        #     error_msg = (
        #         f"Tool '{tool_name}' failed with exception:\n{e}\n\n"
        #         f"Traceback:\n{tb}\n"
        #         "Decide whether to retry, fix the issue, or skip this step."
        #     )
        #     self.logger.error(error_msg)
        #     return {"ok": False, "output": error_msg}


    def _format_tool_usage_ouput(self, id_, tool_name, arguments, output):
        return {
            "tool_call_id": id_,
            "role": "tool",
            "name": tool_name,
            "arguments": arguments,
            "content": output,
        }

    def _create_logs(self):
        def to_dict_safe(msg):
            if isinstance(msg, dict):
                return msg
            elif hasattr(msg, "model_dump"):  # Pydantic/LiteLLM object
                return msg.model_dump()
            else:
                return str(msg)

        safe_messages = [to_dict_safe(m) for m in self.messages]
        created_files = [f for f in os.listdir(self.sandbox_dir) if os.path.isfile(os.path.join(self.sandbox_dir, f))]

        # Build run log
        run_data = {
            "timestamp": utils.time_now(),
            "model": self.model_name,
            "temperature": self.temperature,
            "tools": self.tool_schemas,
            "messages": safe_messages,
            "files_created": created_files,
        }

        utils.append_jsonl(run_data, constants.JSON_LOG_FILE)

    def _call_llm(self, messages):
        # periodically summarize
        total_tokens = 0
        for m in messages:
            total_tokens += self._count_tokens(m.get("content", ""))

        if total_tokens > constants.MAX_CONTEXT_TOKENS and len(messages) > 3:
            messages = self._summarize_history(messages)
            self.messages = messages

        response = completion(
            model=self.model_name,
            temperature=self.temperature,
            supports_system_message=self.model_supports_system_messages,
            messages=messages,
            web_search_options={"search_context_size": "medium"},
            tools=self.tool_schemas,
            tool_choice="auto",
        )
        self.llm_cost += response._hidden_params["response_cost"]

        message = response.choices[0].message

        return message

    def _prompt_llm(self, prompt):
        self.messages.append({"role": "user", "content": prompt})
        return self._call_llm(self.messages)

    def _final_log(self, llm_cost):
        final_logs = {
            "timestamp_final": utils.time_now(),
            "total_completion_cost": llm_cost,
        }

        utils.append_jsonl(final_logs, constants.JSON_LOG_FILE)

    def _process_tool_call(self, tool_call):
        function_name = tool_call.function.name
        raw_args = tool_call.function.arguments

        if not raw_args:
            self.logger.warning(f"Tool '{function_name}' returned empty arguments. Using empty dict.")
            function_args = {}
        else:
            try:
                function_args = json.loads(raw_args)
            except json.JSONDecodeError:
                self.logger.warning(f"Tool '{function_name}' returned invalid JSON: {raw_args}. Using empty dict.")
                function_args = {}

        exec_result = self._safe_execute_tool(function_name, function_args)
        tool_output = utils.truncate_string(exec_result["output"])
        exec_result["output"] = tool_output
        
        self.logger.info(f"Tool result: {tool_output}.")

        tool_message = self._format_tool_usage_ouput(tool_call.id, function_name, function_args, tool_output)
        self.messages.append(tool_message)

        return exec_result

    @abstractmethod
    def _reset_pipeline(self) -> None:
        pass

    @abstractmethod
    def run(self):
        pass