import os
from typing import Dict, Any
import requests
from langchain.llms.openai import OpenAI
from langchain.prompts import PromptTemplate
from langchain.chains import LLMChain

# Simple web search tool using DuckDuckGo instant answer API
class WebSearch:
    def __init__(self):
        self.base = "https://api.duckduckgo.com/"

    def run(self, query: str) -> str:
        params = {
            "q": query,
            "format": "json",
            "no_redirect": 1,
            "no_html": 1,
            "skip_disambig": 1,
        }
        r = requests.get(self.base, params=params)
        data = r.json()
        # Return the abstract text if available
        return data.get("AbstractText", "") or data.get("RelatedTopics", [])[0].get("Text", "")

# Agent that searches and writes a virtual file
class DeepAgent:
    def __init__(self, llm: Any):
        self.llm = llm
        self.search = WebSearch()
        self.virtual_fs: Dict[str, str] = {}

    def run(self, task_description: str) -> None:
        # Step 1: Search the web for relevant info
        search_query = f"{task_description} example"
        context = self.search.run(search_query)
        if not context:
            context = "No context found."

        # Step 2: Ask LLM to generate file content based on context
        prompt = PromptTemplate(
            input_variables=["context", "task_description"],
            template="You are a developer. Based on the following context, write a Python file named output.py that demonstrates the concept described in the task: {task_description}\nContext: {context}\nOutput:\n"
        )
        chain = LLMChain(llm=self.llm, prompt=prompt)
        result = chain.run(context=context, task_description=task_description)

        # Step 3: Store in virtual FS
        self.virtual_fs["output.py"] = result.strip()

    def export_to_real_fs(self, repo_path: str) -> None:
        for filename, content in self.virtual_fs.items():
            full_path = os.path.join(repo_path, filename)
            with open(full_path, "w", encoding="utf-8") as f:
                f.write(content)

# Example usage
if __name__ == "__main__":
    llm = OpenAI(temperature=0.2, model_name="gpt-4o-mini")
    agent = DeepAgent(llm)
    task_desc = "Create a simple Python script that prints 'Hello World'"
    agent.run(task_desc)
    # Export to repository directory
    repo_dir = os.path.abspath(".")
    agent.export_to_real_fs(repo_dir)
    print("Files written:", list(agent.virtual_fs.keys()))