import asyncio, subprocess, re, json
from pathlib import Path
from typing import List, Dict
from model_router import route_chat

class CodingAgent:
    def __init__(self, project_path: str, model: str = "groq-llama-3.3-70b-versatile"):
        self.project_path = Path(project_path)
        self.model = model
        self.status = "idle"
        self.logs = []
        self.pending_changes = []

    async def run_task(self, task: str) -> dict:
        """New agent loop – returns {logs, pending_changes}."""
        self.status = "running"
        self.logs = []
        self.pending_changes = []

        system = f"""You are an AI coding agent. Complete the user's task step by step.
You work in the project folder: {self.project_path}.
Available tools (respond with a JSON object containing "thought", "tool", and "arguments"):
- read_file(filepath)
- write_file(filepath, content)
- edit_file(filepath, search, replace)
- run_command(command)
- list_dir(path)
- task_complete()

Example: {{"thought": "I need to see the code", "tool": "read_file", "arguments": {{"filepath": "main.py"}}}}
Only output the JSON object."""

        messages = [
            {"role": "system", "content": system},
            {"role": "user", "content": task}
        ]

        for step in range(8):
            response = await self._generate(messages)
            self.logs.append(f"Step {step+1}: {response[:150]}...")

            action = self._parse_action(response)
            if not action:
                messages.append({"role": "user", "content": "Invalid JSON. Please output only the JSON action."})
                continue

            thought = action.get("thought", "")
            tool = action.get("tool", "")
            args = action.get("arguments", {})

            result = ""
            if tool == "read_file":
                result = self._read_file(args.get("filepath", ""))
            elif tool == "write_file":
                filepath = args.get("filepath", "")
                content = args.get("content", "")
                self.pending_changes.append({"action": "create", "filepath": filepath, "content": content})
                result = f"Queued creation of {filepath}"
            elif tool == "edit_file":
                filepath = args.get("filepath", "")
                search = args.get("search", "")
                replace = args.get("replace", "")
                current = self._read_file(filepath)
                if search in current:
                    modified = current.replace(search, replace)
                    self.pending_changes.append({"action": "edit", "filepath": filepath, "original": current, "modified": modified})
                    result = f"Queued edit of {filepath}"
                else:
                    result = f"Search string not found in {filepath}"
            elif tool == "run_command":
                result = self._run_command(args.get("command", ""))
            elif tool == "list_dir":
                result = "\n".join(self._list_dir(args.get("path", "")))
            elif tool == "task_complete":
                self.logs.append("Agent marked task complete.")
                break
            else:
                result = f"Unknown tool: {tool}"

            messages.append({"role": "assistant", "content": response})
            messages.append({"role": "user", "content": f"Tool result: {result}"})

        self.status = "awaiting_approval"
        return {"logs": self.logs, "pending_changes": self.pending_changes}

    async def approve_and_apply(self):
        self.status = "applying"
        for change in self.pending_changes:
            try:
                if change["action"] == "create":
                    (self.project_path / change["filepath"]).parent.mkdir(parents=True, exist_ok=True)
                    (self.project_path / change["filepath"]).write_text(change["content"])
                    self.logs.append(f"Created {change['filepath']}")
                elif change["action"] == "edit":
                    (self.project_path / change["filepath"]).write_text(change["modified"])
                    self.logs.append(f"Edited {change['filepath']}")
            except Exception as e:
                self.logs.append(f"Error: {e}")
        self.pending_changes = []
        self.status = "completed"

    def _read_file(self, filepath: str) -> str:
        p = self.project_path / filepath
        return p.read_text() if p.exists() else "File not found."

    def _run_command(self, cmd: str) -> str:
        try:
            r = subprocess.run(cmd, shell=True, cwd=self.project_path, capture_output=True, text=True, timeout=30)
            return r.stdout + r.stderr
        except Exception as e:
            return str(e)

    def _list_dir(self, path: str) -> list:
        p = self.project_path / path if path else self.project_path
        try:
            return [f.name for f in p.iterdir()]
        except:
            return []

    def _parse_action(self, text: str) -> dict | None:
        text = text.strip()
        if text.startswith("```"): text = text[3:]
        if text.endswith("```"): text = text[:-3]
        try:
            return json.loads(text)
        except:
            return None

    async def _generate(self, messages):
        full = ""
        async for token in route_chat(self.model, messages):
            full += token
        return full