import json
import re
from pathlib import Path

from langchain_ollama import ChatOllama

from tools import SearchTool, WriteFileTool
from virtual_fs import VirtualFileSystem

class DeepAgent:
    """A deep‑search agent that can search the web, write virtual files and export them.

    The agent follows a simple plan: the LLM generates a JSON array of steps, each step
    is either a ``search`` or a ``write`` action.  Search results are stored in the
    virtual file system so that later steps can refer to them.
    """

    def __init__(self, llm: ChatOllama):
        self.llm = llm
        self.vfs = VirtualFileSystem()
        self.search_tool = SearchTool()
        self.write_tool = WriteFileTool(self.vfs)

    def _extract_json(self, text: str) -> str:
        """Try to extract a JSON array from an arbitrary string.

        The LLM might prepend or append text.  We look for the first ``[`` and the
        matching ``]`` and slice the string.  If that fails we fall back to the
        original text.
        """
        try:
            start = text.index("[")
            end = text.rindex("]") + 1
            return text[start:end]
        except ValueError:
            return text

    def generate_plan(self, instruction: str) -> list:
        """Ask the LLM to produce a JSON plan for the given instruction.

        Returns a list of step dictionaries.
        """
        prompt = (
            "You are a helpful assistant that can search the web and write files. "
            "Given the instruction below, produce a JSON array of steps. "
            "Each step is an object with \"action\" (\"search\" or \"write\") and "
            "\"params\". For \"search\", params: \"query\". For \"write\", params: "
            "\"filename\" and \"content\". Return only the JSON, no explanation.\n\n"
            f"Instruction: {instruction}"
        )
        response = self.llm.invoke(prompt)
        json_text = self._extract_json(str(response))
        try:
            plan = json.loads(json_text)
        except json.JSONDecodeError as exc:
            raise ValueError(f"Failed to parse plan JSON: {exc}\nResponse: {response}")
        if not isinstance(plan, list):
            raise ValueError("Plan must be a JSON array of steps.")
        return plan

    def execute_plan(self, plan: list):
        for step in plan:
            action = step.get("action")
            params = step.get("params", {})
            if action == "search":
                query = params.get("query")
                if not query:
                    continue
                result = self.search_tool.run(query)
                # Sanitize filename: replace spaces with underscores
                safe_query = re.sub(r"[^a-zA-Z0-9_]+", "_", query)
                self.vfs.write_file(f"search_{safe_query}.txt", result)
            elif action == "write":
                filename = params.get("filename")
                content = params.get("content")
                if not filename or content is None:
                    continue
                self.write_tool.run(filename, content)
            else:
                # Unknown action – skip
                continue

    def run(self, instruction: str, output_dir: str | Path = "./output"):
        """Run the agent for a single instruction.

        Parameters
        ----------
        instruction: str
            The natural‑language instruction for the agent.
        output_dir: str | Path
            Directory where the virtual files will be flushed to disk.
        """
        plan = self.generate_plan(instruction)
        self.execute_plan(plan)
        # Flush virtual FS to disk
        self.vfs.flush_to_disk(Path(output_dir))