"""
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 ast
import inspect
import json
import os
import subprocess
import sys
from pathlib import Path

import requests
from dotenv import load_dotenv

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


# ═════════════════════════════════════════════════════════════════════════════
# PART A – TOOL FUNCTIONS
# ═════════════════════════════════════════════════════════════════════════════
#
# Each tool is a plain Python function decorated with @register_tool.
# The decorator adds the function to TOOL_REGISTRY so the dispatcher
# can call it by name at runtime.

TOOL_REGISTRY: dict[str, callable] = {}


def register_tool(func):
    """Decorator – adds a function to the global tool registry."""
    TOOL_REGISTRY[func.__name__] = func
    return func


@register_tool
def read_file(path: str) -> str:
    """Read a .py or .txt file from the workspace and return its contents."""
    target = (WORKSPACE / path).resolve()
    if not str(target).startswith(str(WORKSPACE.resolve())):
        return "ERROR: path is outside the workspace."
    if not target.exists():
        return f"ERROR: file '{path}' not found."
    if target.suffix not in (".py", ".txt"):
        return f"ERROR: can only read .py and .txt files, got '{target.suffix}'."
    return target.read_text()


@register_tool
def write_file(path: str, content: str) -> str:
    """Write content to a .py or .txt file in the workspace."""
    target = (WORKSPACE / path).resolve()
    if not str(target).startswith(str(WORKSPACE.resolve())):
        return "ERROR: path is outside the workspace."
    if target.suffix not in (".py", ".txt"):
        return f"ERROR: can only write .py and .txt files, got '{target.suffix}'."
    target.parent.mkdir(parents=True, exist_ok=True)
    target.write_text(content)
    return f"OK: wrote {len(content)} chars to {path}."


@register_tool
def list_files(file_glob: str = "*") -> str:
    """List files in the workspace matching the glob pattern."""
    found = sorted(WORKSPACE.glob(file_glob))
    found = [f.relative_to(WORKSPACE) for f in found if f.is_file()]
    if not found:
        return f"No files matching '{file_glob}' in workspace."
    return "\n".join(str(f) for f in found)


@register_tool
def grep_search(pattern: str, file_glob: str = "*.py") -> str:
    """Search for a pattern in workspace files and return matching lines with line numbers."""
    matches = []
    for filepath in sorted(WORKSPACE.glob(file_glob)):
        if filepath.suffix not in (".py", ".txt"):
            continue
        try:
            lines = filepath.read_text().splitlines()
        except Exception:
            continue
        for i, line in enumerate(lines, 1):
            if pattern in line:
                rel = filepath.relative_to(WORKSPACE)
                matches.append(f"{rel}:{i}: {line}")
    if not matches:
        return f"No matches for '{pattern}' in {file_glob}."
    return "\n".join(matches)


@register_tool
def run_python(path: str) -> str:
    """Execute a Python file in the workspace and return stdout and stderr."""
    target = (WORKSPACE / path).resolve()
    if not str(target).startswith(str(WORKSPACE.resolve())):
        return "ERROR: path is outside the workspace."
    if not target.exists():
        return f"ERROR: file '{path}' not found."
    result = subprocess.run(
        [sys.executable, str(target)],
        capture_output=True, text=True, timeout=30,
        cwd=str(WORKSPACE),
    )
    output = ""
    if result.stdout:
        output += f"STDOUT:\n{result.stdout}"
    if result.stderr:
        output += f"STDERR:\n{result.stderr}"
    output += f"\nExit code: {result.returncode}"
    return output.strip()


@register_tool
def validate_python(path: str) -> str:
    """Check whether a Python file has valid syntax using ast.parse."""
    target = (WORKSPACE / path).resolve()
    if not str(target).startswith(str(WORKSPACE.resolve())):
        return "ERROR: path is outside the workspace."
    if not target.exists():
        return f"ERROR: file '{path}' not found."
    source = target.read_text()
    try:
        ast.parse(source)
        return "OK: syntax is valid."
    except SyntaxError as e:
        return f"SYNTAX ERROR: {e}"


@register_tool
def done(summary: str) -> str:
    """Signal that the agent has finished its task."""
    return f"DONE: {summary}"


# ═════════════════════════════════════════════════════════════════════════════
# PART B – TOOL DISPATCHER
# ═════════════════════════════════════════════════════════════════════════════

def build_tool_description() -> str:
    """Auto-generate tool descriptions from function signatures and docstrings."""
    lines = []
    for name, func in TOOL_REGISTRY.items():
        sig = inspect.signature(func)
        params = []
        for pname, param in sig.parameters.items():
            if param.default is inspect.Parameter.empty:
                params.append(f'"{pname}": "<value>"')
            else:
                params.append(f'"{pname}": "<optional, default={param.default!r}>"')
        param_str = ", ".join(params)
        doc = (func.__doc__ or "").strip().split("\n")[0]
        lines.append(f"  - {name}({{{param_str}}}): {doc}")
    return "\n".join(lines)


def dispatch_tool(tool_name: str, arguments: dict) -> str:
    """Call a tool by name with the given arguments."""
    if tool_name not in TOOL_REGISTRY:
        return f"ERROR: unknown tool '{tool_name}'. Available: {list(TOOL_REGISTRY.keys())}"
    func = TOOL_REGISTRY[tool_name]
    try:
        return func(**arguments)
    except TypeError as e:
        return f"ERROR calling {tool_name}: {e}"
    except Exception as e:
        return f"ERROR in {tool_name}: {type(e).__name__}: {e}"


# ═════════════════════════════════════════════════════════════════════════════
# PART C – 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.

<capabilities>
You can:
  - Read .py and .txt files from the workspace
  - Write .py and .txt files to the workspace
  - List files in the workspace
  - Search for patterns in files using grep
  - Execute Python files and see their output
  - Validate Python syntax using ast.parse
  - Signal completion when the task is done
</capabilities>

<tools>
{build_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>
"""


# ═════════════════════════════════════════════════════════════════════════════
# PART D – CODING AGENT CLASS
# ═════════════════════════════════════════════════════════════════════════════

def truncate_result(result: str) -> str:
    """Truncate a tool result that exceeds MAX_RESULT_LENGTH."""
    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:
    """Drop old messages when history exceeds MAX_HISTORY_CHARS.
    Always keeps the system prompt (index 0) and original task (index 1).
    """
    total = sum(len(m["content"]) for m in messages)
    if total <= MAX_HISTORY_CHARS:
        return messages

    head = messages[:2]
    tail = messages[2:]
    original_task = messages[1]["content"] if len(messages) > 1 else ""

    while tail and sum(len(m["content"]) for m in head + tail) > MAX_HISTORY_CHARS:
        tail.pop(0)

    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 _strip_code_fences(text: str) -> str:
    """Remove markdown code fences (```json ... ```) from a string."""
    text = text.strip()
    if text.startswith("```"):
        lines = text.split("\n")
        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,
            "temperature": 0.2,
            "max_tokens": 4096,
            "stream": False,
        }

        response = requests.post(self.api_url, headers=headers, json=payload, timeout=60)

        if response.status_code != 200:
            raise Exception(f"API Error {response.status_code}: {response.text}")

        data = response.json()
        if "choices" in data and len(data["choices"]) > 0:
            return data["choices"][0]["message"]["content"]
        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

    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)
            action = json.loads(raw)
        except json.JSONDecodeError:
            action = {
                "thought": "Could not parse LLM response as JSON.",
                "tool": "done",
                "arguments": {"summary": "Stopped: JSON parse error."},
            }
            raw = json.dumps(action)
        except Exception as e:
            action = {
                "thought": f"API call failed: {e}",
                "tool": "done",
                "arguments": {"summary": f"Stopped: {e}"},
            }
            raw = json.dumps(action)

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

    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

        # 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 = dispatch_tool(tool_name, arguments)
        result = truncate_result(result)

        # Build feedback – nudge agent to replan on errors
        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>"
            )

        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>"
            ),
        })

    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