"""Basic agent class. See https://mini-swe-agent.com/latest/advanced/control_flow/ for visual explanation."""

import re
import subprocess
import time

from jinja2 import StrictUndefined, Template
from pydantic import BaseModel

from cooperbench.agents.mini_swe_agent import Environment, Model
from cooperbench.agents.mini_swe_agent.connectors.messaging import MessagingConnector
from cooperbench.agents.mini_swe_agent.utils.log import logger


class AgentConfig(BaseModel):
    # Check the config files in cooperbench/agents/mini_swe_agent/config for example settings
    system_template: str
    instance_template: str
    timeout_template: str
    format_error_template: str
    action_observation_template: str
    action_regex: str = r"```bash\s*\n(.*?)\n```"
    step_limit: int = 50
    cost_limit: float = 1.0


class NonTerminatingException(Exception):
    """Raised for conditions that can be handled by the agent."""


class FormatError(NonTerminatingException):
    """Raised when the LM's output is not in the expected format."""


class ExecutionTimeoutError(NonTerminatingException):
    """Raised when the action execution timed out."""


class TerminatingException(Exception):
    """Raised for conditions that terminate the agent."""


class Submitted(TerminatingException):
    """Raised when the LM declares that the agent has finished its task."""


class LimitsExceeded(TerminatingException):
    """Raised when the agent has reached its cost or step limit."""


class DefaultAgent:
    def __init__(
        self,
        model: Model,
        env: Environment,
        *,
        comm: "MessagingConnector | None" = None,
        agent_id: str = "agent",
        config_class: type = AgentConfig,
        **kwargs,
    ):
        self.config = config_class(**kwargs)
        self.messages: list[dict] = []
        self.model = model
        self.env = env
        self.comm = comm  # Optional inter-agent communication
        self.agent_id = agent_id
        self.step_num = 0
        self.extra_template_vars = {}

    def log(self, msg: str):
        """Log message with agent prefix."""
        logger.debug(f"[{self.agent_id}] {msg}")

    def render_template(self, template: str, **kwargs) -> str:
        template_vars = self.config.model_dump() | self.env.get_template_vars() | self.model.get_template_vars()
        return Template(template, undefined=StrictUndefined).render(
            **kwargs, **template_vars, **self.extra_template_vars
        )

    def add_message(self, role: str, content: str, **kwargs):
        self.messages.append({"role": role, "content": content, "timestamp": time.time(), **kwargs})

    def run(self, task: str, **kwargs) -> tuple[str, str]:
        """Run step() until agent is finished. Return exit status & message"""
        self.extra_template_vars |= {"task": task, **kwargs}
        self.messages = []
        self.add_message("system", self.render_template(self.config.system_template))
        self.add_message("user", self.render_template(self.config.instance_template))
        self.step_num = 0
        while True:
            self.step_num += 1
            self.log(f"Step {self.step_num}")
            try:
                self.step()
            except NonTerminatingException as e:
                self.log(f"Error: {type(e).__name__}")
                self.add_message("user", str(e))
            except TerminatingException as e:
                self.log(f"Finished: {type(e).__name__}")
                self.add_message("user", str(e))
                return type(e).__name__, str(e)

    def step(self) -> dict:
        """Query the LM, execute the action, return the observation."""
        # Check for inter-agent messages before querying LLM
        if self.comm:
            messages = self.comm.receive()
            for msg in messages:
                ts = msg.get("timestamp", "")[:19].replace("T", " ")
                self.log(f"INBOX: [{msg['from']} @ {ts}] {msg['content']}")
                self.add_message("user", f"[Message from {msg['from']}]: {msg['content']}")
        return self.get_observation(self.query())

    def query(self) -> dict:
        """Query the model and return the response."""
        if 0 < self.config.step_limit <= self.model.n_calls or 0 < self.config.cost_limit <= self.model.cost:
            raise LimitsExceeded()
        logger.debug(f"Querying LLM (call #{self.model.n_calls + 1})...")
        response = self.model.query(self.messages)
        logger.debug(f"Got response ({len(response.get('content', ''))} chars)")
        self.add_message("assistant", **response)
        return response

    def get_observation(self, response: dict) -> dict:
        """Execute the action and return the observation."""
        output = self.execute_action(self.parse_action(response))
        observation = self.render_template(self.config.action_observation_template, output=output)
        self.add_message("user", observation)
        return output

    def parse_action(self, response: dict) -> dict:
        """Parse the action from the message. Returns the action."""
        actions = re.findall(self.config.action_regex, response["content"], re.DOTALL)
        if len(actions) == 1:
            return {"action": actions[0].strip(), **response}
        raise FormatError(self.render_template(self.config.format_error_template, actions=actions))

    def execute_action(self, action: dict) -> dict:
        cmd = action["action"]
        self.log(f"> {cmd}")

        # Handle send_message command(s) for inter-agent communication
        # Detect ALL send_message calls anywhere in the command (handles compound commands)
        if self.comm:
            # Find all send_message calls (double and single quotes)
            matches = re.findall(r'send_message\s+(\w+)\s+"([^"]*)"', cmd)
            matches += re.findall(r"send_message\s+(\w+)\s+'([^']*)'", cmd)

            if matches:
                # Send all messages
                outputs = []
                for recipient, content in matches:
                    result = self._handle_send_message(recipient, content)
                    outputs.append(result["output"])
                    self.log(f"(exit {result['returncode']})")
                    self.log(result["output"])

                # Strip ALL send_message calls and execute remaining command
                remaining_cmd = self._strip_send_message(cmd)
                combined_output = "\n".join(outputs)

                if not remaining_cmd.strip():
                    return {"output": combined_output, "returncode": 0, "action": cmd}

                # Execute remaining command and combine outputs
                self.log(f"> {remaining_cmd}")
                try:
                    env_output = self.env.execute(remaining_cmd)
                    combined_output += "\n" + env_output.get("output", "")
                    output = {"output": combined_output, "returncode": env_output.get("returncode", 0), "action": cmd}
                except Exception as e:
                    output = {
                        "output": combined_output + f"\n(remaining command failed: {e})",
                        "returncode": 1,
                        "action": cmd,
                    }

                self.log(f"(exit {output.get('returncode', '?')})")
                self.log(output.get("output", ""))
                self.has_finished(output)
                return output

        try:
            output = self.env.execute(cmd)
        except (TimeoutError, subprocess.TimeoutExpired) as e:
            output_str = e.output.decode("utf-8", errors="replace") if getattr(e, "output", None) else ""
            self.log("(timeout)")
            self.log(output_str)
            raise ExecutionTimeoutError(
                self.render_template(self.config.timeout_template, action=action, output=output_str)
            )
        except Exception as e:
            # Handle unexpected errors (e.g., encoding issues) gracefully
            self.log(f"(error: {e})")
            output = {"output": f"Command execution error: {e}", "returncode": 1, "action": cmd}
            return output

        self.log(f"(exit {output.get('returncode', '?')})")
        self.log(output.get("output", ""))
        self.has_finished(output)
        return output | {"action": cmd}

    def _strip_send_message(self, cmd: str) -> str:
        """Remove send_message call from a compound command."""
        # Remove: send_message agent "msg"
        cmd = re.sub(r'send_message\s+\w+\s+"[^"]*"', "", cmd)
        cmd = re.sub(r"send_message\s+\w+\s+'[^']*'", "", cmd)
        # Clean up leftover && or ||
        cmd = re.sub(r"^\s*&&\s*", "", cmd)
        cmd = re.sub(r"\s*&&\s*$", "", cmd)
        cmd = re.sub(r"&&\s*&&", "&&", cmd)
        cmd = re.sub(r"\|\|\s*\|\|", "||", cmd)
        return cmd.strip()

    def _handle_send_message(self, recipient: str, content: str) -> dict:
        """Handle send_message command."""
        self.comm.send(recipient, content)
        logger.debug(f"Sent message to {recipient}: {content[:50]}...")
        return {"output": f"Message sent to {recipient}", "returncode": 0}

    def has_finished(self, output: dict[str, str]):
        """Raises Submitted exception with final output if the agent has finished its task."""
        lines = output.get("output", "").lstrip().splitlines(keepends=True)
        if lines and lines[0].strip() in ["MINI_SWE_AGENT_FINAL_OUTPUT", "COMPLETE_TASK_AND_SUBMIT_FINAL_OUTPUT"]:
            raise Submitted("".join(lines[1:]))
