"""
Coding Agent
============
Autonomous AI coding agent based on the Plan→Act→Observe→Fix→Done loop.
Structure follows ex05_coding_agent_solution.py from the course.

Instead of a blocking CLI loop (input()), the CodingAgent class exposes
step-by-step methods so Streamlit can drive the loop via session_state:

    agent.start_task(task)          # initialise
    action = agent.propose_next_action()  # ask LLM → returns action, does NOT execute
    result = agent.approve()        # execute pending action
    agent.reject(feedback)          # skip action, inject user feedback
"""

import json
import os
import re
from pathlib import Path
import asyncio

import requests
from dotenv import load_dotenv
from backend.agent.mcp_server_adapter import MCPToolAdapter

from backend.managers.debug_logger import get_logger
logger = get_logger(__name__)

# ── mcp server initialization ────────────────────────────────────────────────────────────────
adapter = MCPToolAdapter()
logger.info("MCPToolAdapter created. Listing all tools from servers...")
asyncio.run(adapter.initialize_all_servers())
logger.info("Listed tools from all servers")

load_dotenv()

# ── Workspace ────────────────────────────────────────────────────────────────
# Two levels up from backend/agent/ → project root → workspace/
WORKSPACE = Path(__file__).resolve().parents[2] / "workspace"
WORKSPACE.mkdir(exist_ok=True)

# ── Agent limits ─────────────────────────────────────────────────────────────
MAX_ITERATIONS    = 100
MAX_RESULT_LENGTH = 10_000
MAX_HISTORY_CHARS = 80_000


# ═════════════════════════════════════════════════════════════════════════════
# Tool dispatching and result handling
# ═════════════════════════════════════════════════════════════════════════════

def build_all_tool_description() -> str:
    """Build a formatted string listing every registered MCP tool.

    The returned string is embedded verbatim in the SYSTEM_PROMPT so the LLM
    knows which tools exist and what arguments they expect.

    Returns:
        Newline-separated list of tool descriptions in the format
        ``"- <tool_name>: <description>"``.
    """
    all_tools = adapter.get_all_tools()
    logger.info("Building tool description for %s tools.", str(len(all_tools)))

    descriptions = []
    for tool in all_tools:
        descriptions.append(f"- {tool['tool_name']}: {tool['tool_description']}")

    return "\n".join(descriptions)

async def dispatch_tool(tool_name: str, arguments: dict) -> str:
    """Execute a named tool and return its output as a plain string.

    Handles the special "done" pseudo-tool locally (it signals completion and
    is never forwarded to an MCP server).  All other tools are forwarded to the
    MCPToolAdapter which routes them to the correct MCP server process.

    Args:
        tool_name: Name of the tool to execute (e.g. "write_file", "done").
        arguments: Dict of arguments for the tool.

    Returns:
        The tool's text output, a "DONE: ..." completion message, or an error
        string beginning with "Tool error:" / "Error calling tool:" on failure.
    """
    if tool_name == "done":
        # The "done" tool is a sentinel — it lives only in the agent protocol,
        # not in any MCP server, so we resolve it directly here.
        summary = arguments.get("summary", "Task completed.")
        return f"DONE: {summary}"

    try:
        logger.info("Calling tool '%s' in dispatch_tool through MCPToolAdapter...", tool_name)
        result = await adapter.call_tool(tool_name, arguments)

        logger.info("Result from tool '%s' received", tool_name)

        if result.isError:
            # MCP servers signal tool-level errors via the isError flag rather
            # than raising exceptions, so we surface them explicitly.
            texts = [block.text for block in result.content if block.type == "text"]
            logger.warning("Result from '%s' is Error", tool_name)
            return f"Tool error: {' '.join(texts)}"

        texts = [block.text for block in result.content if block.type == "text"]
        return "\n".join(texts)

    except Exception as e:
        logger.exception("Error calling tool '%s' with argument: %s", tool_name, arguments)
        return f"Error calling tool '{tool_name}': {e}"

# ═════════════════════════════════════════════════════════════════════════════
# SYSTEM PROMPT
# ═════════════════════════════════════════════════════════════════════════════

SYSTEM_PROMPT = f"""\
You are a coding agent that helps users with Python programming tasks.
You work inside a workspace directory and have access to tools.
You solve tasks by interacting with a workspace through a
dynamic set of tools via the Model Context Protocol (MCP).

<capabilities> 
You have access to a workspace where you can manage files,
analyze code, and execute Python scripts.
You can call tools to interact with the workspace and get feedback.
You can write and read files, list directory contents, search for patterns,
validate Python syntax, and run Python code.
You can access web search and page fetching tools to gather information from the internet.
You can use these capabilities to iteratively work towards completing the user's task.
</capabilities>

<tools>
{build_all_tool_description()}
</tools>


<workflow>
For every user request, follow this workflow:
  1. PLAN:    Think about what steps are needed. List them in "thought".
  2. ACT:     Choose ONE tool to call for the current step.
  3. OBSERVE: You will receive the tool output. Analyse it carefully.
  4. REPLAN:  If the result was unexpected, revise your plan in "thought".
  5. REPEAT:  Go back to step 2 if more work is needed.
  6. DONE:    Call the "done" tool when the task is complete.
</workflow>

<response_format>
You MUST respond with a JSON object every time:
{{
  "thought": "<your reasoning about what to do next>",
  "tool": "<tool name>",
  "arguments": {{ <tool arguments> }}
}}

Example:
{{
  "thought": "I need to write fibonacci.py first, then validate and run it.",
  "tool": "write_file",
  "arguments": {{"path": "fibonacci.py", "content": "def fib(n): ..."}}
}}
</response_format>

<rules>
  - Always plan before acting.
  - Call exactly ONE tool per response.
  - After writing code, ALWAYS validate it with validate_python.
  - After validation passes, run it with run_python to verify correctness.
  - If an error occurs, analyse it and try to fix it (up to 3 retries).
  - Stay within the workspace directory.
  - When the task is fully complete, call the "done" tool.
  - If you receive a <human_message>, acknowledge it and adjust your plan.
  - If you receive a <replan> tag, revise your plan before choosing the next tool.
</rules>
"""


# ═════════════════════════════════════════════════════════════════════════════
# CODING AGENT CLASS
# ═════════════════════════════════════════════════════════════════════════════

def truncate_result(result: str) -> str:
    """Truncate a tool result that exceeds MAX_RESULT_LENGTH."""
    logger.info("Result has been truncated")
    if len(result) <= MAX_RESULT_LENGTH:
        return result
    half = MAX_RESULT_LENGTH // 2
    return (
        result[:half]
        + f"\n\n... [TRUNCATED – {len(result)} chars total] ...\n\n"
        + result[-half:]
    )


def trim_messages(messages: list) -> list:
    """Kürzt die Konversations-History wenn sie das Kontextfenster überschreitet.

    Behält immer den System-Prompt (Index 0) und die ursprüngliche Aufgabe (Index 1).
    Entfernt die ältesten Nachrichten zuerst und injiziert danach einen Erinnerungs-
    Hinweis damit der Agent den Überblick behält.

    Args:
        messages: Vollständige Konversations-History als Liste von {role, content} Dicts.

    Returns:
        Gekürzte History mit maximal MAX_HISTORY_CHARS Zeichen, immer mit Head + Reminder + Tail.
    """
    logger.info("Message is being trimmed")

    total = sum(len(m["content"]) for m in messages)
    if total <= MAX_HISTORY_CHARS:
        return messages

    # Protect the two anchor messages that must never be discarded.
    head = messages[:2]
    tail = messages[2:]
    original_task = messages[1]["content"] if len(messages) > 1 else ""

    # Drop the oldest non-anchor messages first until we are under the limit.
    while tail and sum(len(m["content"]) for m in head + tail) > MAX_HISTORY_CHARS:
        tail.pop(0)

    # After trimming, inject a reminder so the agent doesn't lose track of its goal.
    # Without this the agent might restart the task or repeat work it already did.
    reminder = {
        "role": "user",
        "content": (
            "<system_note>Earlier conversation history was trimmed to fit the context window. "
            f"REMINDER – your original task was:\n{original_task}\n"
            "Continue working towards completing this task. "
            "Do NOT start over or redo work already completed.</system_note>"
        ),
    }
    return head + [reminder] + tail

def _repair_json_strings(text: str) -> str:
    """Replace unescaped control characters inside JSON string values.

    LLMs frequently emit literal newlines, tabs, or carriage-returns inside
    long string values (e.g. code content), which is invalid JSON.  This
    function replaces those characters with their proper ``\\n`` / ``\\t`` /
    ``\\r`` escape sequences without touching structural whitespace that lives
    outside of string literals.

    The parser is a simple state-machine that tracks whether the current
    character is inside a quoted string, respecting backslash escapes.

    Args:
        text: Raw JSON text that may contain unescaped control characters.

    Returns:
        Repaired JSON text with control characters properly escaped inside strings.
    """
    result: list[str] = []
    in_string = False
    escape = False
    _escapes = {'\n': '\\n', '\r': '\\r', '\t': '\\t'}
    for ch in text:
        if escape:
            # The previous character was a backslash — emit this char literally
            # and reset the escape flag.
            result.append(ch)
            escape = False
            continue
        if ch == '\\' and in_string:
            result.append(ch)
            escape = True
            continue
        if ch == '"':
            # Toggle string-mode on every unescaped double quote.
            in_string = not in_string
            result.append(ch)
            continue
        if in_string and ch in _escapes:
            # Replace the bare control character with its escape sequence.
            result.append(_escapes[ch])
            continue
        result.append(ch)
    return ''.join(result)


def extract_json(text: str) -> str:
    """
    Extract and repair a JSON object or array from an LLM response that may
    contain extra prose, markdown code fences, or unescaped control characters.

    Strategy:
      1. Strip markdown ```json ... ``` or ``` ... ``` fences.
      2. Find the first '{' or '[' and extract to the matching closing bracket.
      3. Repair unescaped newlines/tabs inside string values.

    Returns the cleaned JSON string, or the original text as a fallback
    (so json.loads can raise a meaningful error with context).
    """

    if text is None:
        return "" 

    # 1. Strip markdown fences
    fenced = re.sub(r"```(?:json)?\s*([\s\S]*?)\s*```", r"\1", text.strip())
    if fenced != text.strip():
        return _repair_json_strings(fenced.strip())

    # 2. Find first JSON container and extract to matching close
    extracted = text
    for start_char, end_char in [('{', '}'), ('[', ']')]:
        idx = text.find(start_char)
        if idx == -1:
            continue
        depth = 0
        in_string = False
        escape = False
        for i, ch in enumerate(text[idx:], start=idx):
            if escape:
                escape = False
                continue
            if ch == '\\' and in_string:
                escape = True
                continue
            if ch == '"':
                in_string = not in_string
                continue
            if in_string:
                continue
            if ch == start_char:
                depth += 1
            elif ch == end_char:
                depth -= 1
                if depth == 0:
                    extracted = text[idx: i + 1]
                    break
        break

    # 3. Repair unescaped control characters inside string values
    return _repair_json_strings(extracted)


def _strip_code_fences(text: str) -> str:
    """Remove a single wrapping markdown code fence from a string.

    Handles both `` ```json `` and plain `` ``` `` opening fences.  If the text
    does not start with a fence the string is returned unchanged.

    Args:
        text: Raw LLM response that may be wrapped in a markdown code block.

    Returns:
        The text with the opening fence line and optional closing `` ``` `` line
        removed, stripped of surrounding whitespace.
    """
    if text is None:
        return ""

    text = text.strip()
    if text.startswith("```"):
        lines = text.split("\n")
        # Omit the last line only if it is a closing fence; otherwise keep everything.
        end = -1 if lines[-1].strip() == "```" else len(lines)
        text = "\n".join(lines[1:end])
    return text.strip()


class CodingAgent:
    """Step-by-step coding agent for use in Streamlit.

    Usage:
        agent = CodingAgent()
        agent.start_task("Write a fibonacci function")

        action = agent.propose_next_action()
        # → {"thought": "...", "tool": "write_file", "arguments": {...}}
        # Show action in UI, wait for user input

        result = agent.approve()
        # → {"tool": "write_file", "result": "OK: ...", "is_done": False}

        agent.reject("Please use recursion instead")
        # → agent replans on next propose_next_action()
    """

    def __init__(self):
        self.messages: list = []
        self.pending_action: dict | None = None
        self.is_done: bool = False
        self.iteration: int = 0
        self._setup_api()

    def _setup_api(self):
        """Load API configuration from environment variables."""
        self.api_url = f"http://{os.getenv('HOST')}:{os.getenv('PORT')}/v1/chat/completions"
        self.api_key = os.getenv("API_KEY")
        self.model   = os.getenv("MODEL")

    def _call_api(self, messages: list) -> str:
        """Make a raw API call and return the response content string."""
        
        headers = {"Content-Type": "application/json"}
        if self.api_key and self.api_key != "EMPTY":
            headers["Authorization"] = f"Bearer {self.api_key}"

        payload = {
            "model": self.model,
            "messages": messages,
            # Low temperature keeps the agent's tool selections deterministic and
            # reduces the chance of hallucinated tool names or argument formats.
            "temperature": 0.2,
            "max_tokens": 4096,
            "stream": False,
        }

        try:
            response = requests.post(
                self.api_url, 
                headers=headers, 
                json=payload, 
                timeout=60)
            response.raise_for_status()
            logger.info("LLM API response requested")
        except requests.RequestException as exc:
            logger.exception("API Error; HTTP-Fehler: %s", exc)
            raise Exception(f"HTTP-Fehler: {exc}") from exc
        
        data = response.json()
        if "choices" in data and len(data["choices"]) > 0:
            logger.info("valid API output, data returned")
            return data["choices"][0]["message"]["content"]
        
        logger.error("Invalid API response format")
        raise Exception("Invalid API response format")
    

    # ── Public interface ──────────────────────────────────────────────────────

    def start_task(self, task: str) -> None:
        """Initialise the agent with a new task. Resets all state."""
        self.messages = [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user",   "content": task},
        ]
        self.pending_action = None
        self.is_done        = False
        self.iteration      = 0
        logger.info("New ask initialized")

    async def propose_next_action(self) -> dict:
        """Ask the LLM what to do next.

        Returns the parsed action dict without executing anything.
        The action is stored internally as pending_action until
        approve() or reject() is called.

        Returns:
            {"thought": str, "tool": str, "arguments": dict}
        """
        if self.is_done:
            return {"thought": "Task already completed.", "tool": "done", "arguments": {}}

        if self.iteration >= MAX_ITERATIONS:
            return {"thought": "Max iterations reached.", "tool": "done",
                    "arguments": {"summary": "Stopped: max iterations reached."}}
        
        self.iteration += 1
        self.messages = trim_messages(self.messages)

        try:
            raw = self._call_api(self.messages)
            raw = _strip_code_fences(raw)
            cleaned = extract_json(raw)
            action = json.loads(cleaned)
            logger.info("Propose next action successfull")
        except json.JSONDecodeError:
            action = {
                "thought": "Could not parse LLM response as JSON.",
                "tool": "done",
                "arguments": {"summary": "Stopped: JSON parse error."},
            }
            raw = json.dumps(action)
            logger.critical("Parsing API response into valid JASON failed in Step 'propose_next_action'")
        except Exception as e:
            action = {
                "thought": f"API call failed: {e}",
                "tool": "done",
                "arguments": {"summary": f"Stopped: {e}"},
            }
            raw = json.dumps(action)
            logger.critical("API call faliled in Step %s: %s", self.iteration, e)

        self.pending_action = {"raw": raw, "action": action}
        return action

    async def approve(self) -> dict:
        """Execute the pending action and return the result.

        Returns:
            {"tool": str, "arguments": dict, "result": str, "is_done": bool}
        """
        if not self.pending_action:
            raise Exception("No pending action. Call propose_next_action() first.")

        raw       = self.pending_action["raw"]
        action    = self.pending_action["action"]
        tool_name = action.get("tool", "")
        arguments = action.get("arguments", {})

        # Append the assistant message to history
        self.messages.append({"role": "assistant", "content": raw})
        self.pending_action = None

        logger.info("Messages prepared after approval")

        # Handle completion
        if tool_name == "done":
            self.is_done = True
            return {
                "tool":      "done",
                "arguments": arguments,
                "result":    arguments.get("summary", "Task completed."),
                "is_done":   True,
            }

        # Execute the tool
        result = await dispatch_tool(tool_name, arguments)
        result = truncate_result(result)
        logger.info("Tool called and result truncated")

        # Wrap the tool output in an XML tag so the LLM can easily find it.
        # Append a <replan> tag on errors to force the agent to reconsider
        # its plan rather than blindly retrying the same failing action.
        feedback = f'<tool_result tool="{tool_name}">\n{result}\n</tool_result>'
        if result.startswith("ERROR") or result.startswith("SYNTAX ERROR"):
            feedback += (
                "\n\n<replan>The tool returned an error. "
                "Re-examine your plan: what went wrong and what should you do differently? "
                "State your revised plan in your next thought.</replan>"
            )
            logger.warning("Error Message in the tool result, replan-feedback will appended")

        self.messages.append({"role": "user", "content": feedback})

        return {
            "tool":      tool_name,
            "arguments": arguments,
            "result":    result,
            "is_done":   False,
        }

    def follow_up(self, question: str) -> None:
        """Continue a completed task with a follow-up question.

        Resets the done-flag and appends the user's question to the existing
        conversation, so the agent retains full context of what was already done.
        Call propose_next_action() afterwards to continue the loop.
        """
        self.is_done = False
        self.messages.append({
            "role": "user",
            "content": (
                f"<human_message>{question}</human_message>\n"
                "<replan>The user has a follow-up question or correction regarding "
                "the task you just completed. Review what you already did and "
                "address their question accordingly.</replan>"
            ),
        })
        logger.info("Follow-up message appended.")

    def reject(self, feedback: str) -> None:
        """Reject the pending action and inject user feedback.

        The pending action is NOT executed. On the next call to
        propose_next_action() the agent will replan based on the feedback.
        """
        if not self.pending_action:
            return

        raw       = self.pending_action["raw"]
        action    = self.pending_action["action"]
        tool_name = action.get("tool", "unknown")

        self.messages.append({"role": "assistant", "content": raw})
        self.messages.append({
            "role": "user",
            "content": (
                f"<human_message>{feedback}</human_message>\n"
                "<replan>The user has given you guidance BEFORE you executed "
                f"your proposed action ({tool_name}). Do NOT proceed with that action. "
                "Revise your plan to incorporate their feedback and state your "
                "updated plan in your next thought.</replan>"
            ),
        })
        self.pending_action = None
        logger.info("Rejection message appended.")