import os
import subprocess
import re
import shutil

WORKSPACE_DIR = os.path.abspath("workspace")

class AgentWorkspace:
    def __init__(self):
        if not os.path.exists(WORKSPACE_DIR):
            os.makedirs(WORKSPACE_DIR, exist_ok=True)
        self.root = WORKSPACE_DIR

    def list_files(self):
        return sorted(os.listdir(self.root))

    def save_file(self, filename, content):
        # Sanitize filename
        filename = os.path.basename(filename)
        path = os.path.join(self.root, filename)
        with open(path, "w", encoding="utf-8") as f:
            f.write(content)
        return path

    def read_file(self, filename):
        filename = os.path.basename(filename)
        path = os.path.join(self.root, filename)
        if os.path.exists(path):
            with open(path, "r", encoding="utf-8") as f:
                return f.read()
        return None

    def delete_file(self, filename):
        filename = os.path.basename(filename)
        path = os.path.join(self.root, filename)
        if os.path.exists(path):
            os.remove(path)
            return True
        return False

class CodeExecutor:
    def __init__(self, workspace):
        self.workspace = workspace

    def run_command(self, command):
        """Runs a shell command in the workspace directory."""
        try:
            # Basic security check: prevent obvious malicious commands
            forbidden = ["rm -rf /", ":(){ :|:& };:"]
            if any(f in command for f in forbidden):
                return "Security Warning: Command blocked."

            result = subprocess.run(
                command,
                shell=True,
                capture_output=True,
                text=True,
                cwd=self.workspace.root,
                timeout=30  # 30 seconds timeout
            )
            output = result.stdout
            if result.stderr:
                output += f"\n[STDERR]\n{result.stderr}"
            return output if output.strip() else "[No Output]"
        except subprocess.TimeoutExpired:
            return "Error: Command timed out after 30 seconds."
        except Exception as e:
            return f"Error executing command: {str(e)}"

    def run_command_safe(self, command):
        """Runs a shell command and returns structured result (stdout, stderr, returncode)."""
        try:
            # Basic security check
            forbidden = ["rm -rf /", ":(){ :|:& };:"]
            if any(f in command for f in forbidden):
                return "", "Security Warning: Command blocked.", 1

            result = subprocess.run(
                command,
                shell=True,
                capture_output=True,
                text=True,
                cwd=self.workspace.root,
                timeout=30
            )
            return result.stdout, result.stderr, result.returncode
        except subprocess.TimeoutExpired:
            return "", "Error: Command timed out after 30 seconds.", 124
        except Exception as e:
            return "", f"Error executing command: {str(e)}", 1

    def run_python(self, filename):
        return self.run_command(f"python3 {filename}")

    def run_python_safe(self, filename):
        return self.run_command_safe(f"python3 {filename}")

def extract_code_blocks(text):
    """Extracts code blocks from markdown text. Returns list of dicts."""
    # Pattern to find ```language\ncode\n```
    # Also handles cases where language might be missing
    # Updated to handle truncated blocks (missing closing backticks)
    pattern = r"```(\w*)\n(.*?)(?:```|$)"
    matches = re.findall(pattern, text, re.DOTALL)
    blocks = []
    for lang, code in matches:
        # cleanup
        lang = lang.strip().lower()
        code = code.strip()
        if code:
            blocks.append({"language": lang if lang else "text", "code": code})
    return blocks

def parse_filename_from_code(code_block):
    """
    Attempts to guess the filename from the first line of code (e.g. # filename: script.py)
    """
    first_line = code_block.split('\n')[0].strip()
    if first_line.startswith(("#", "//")):
        # Look for filename: ... or just the name
        match = re.search(r'filename:\s*([\w\-\.]+)', first_line, re.IGNORECASE)
        if match:
            return match.group(1)
    return None