import os
import asyncio
import shutil
import json
import re
from pathlib import Path
from typing import List

import httpx
from bs4 import BeautifulSoup

from langchain_openai import ChatOpenAI, OpenAIEmbeddings
from langchain_core.messages import HumanMessage
from langchain.tools import tool
from deepagents import create_deep_agent
from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend

# ---------- Configuration ----------
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
if not OPENAI_API_KEY:
    raise EnvironmentError("Please set the OPENAI_API_KEY environment variable.")

LLM = ChatOpenAI(
    model="openai/gpt-oss-20b:free",
    base_url="https://openrouter.ai/api/v1",
    api_key=OPENAI_API_KEY,
    temperature=0.0,
)

# Workspace where the virtual filesystem lives
WORKSPACE_DIR = Path("./workspace")
OUTPUT_DIR = Path("./output")
WORKSPACE_DIR.mkdir(exist_ok=True)
OUTPUT_DIR.mkdir(exist_ok=True)

backend = CompositeBackend(
    [
        LocalShellBackend(workspace_dir=str(WORKSPACE_DIR)),
        FilesystemBackend(),
    ]
)

# ---------- Tools ----------
@tool
def web_search(query: str) -> str:
    """
    Perform a simple web search using DuckDuckGo and return the titles and URLs of the top 3 results.
    The result is a JSON string with a list of objects: [{"title": "...", "url": "..."}].
    """
    try:
        url = "https://html.duckduckgo.com/html/"
        params = {"q": query}
        headers = {"User-Agent": "Mozilla/5.0"}
        resp = httpx.get(url, params=params, headers=headers, timeout=15.0)
        resp.raise_for_status()
        soup = BeautifulSoup(resp.text, "html.parser")
        results = []
        for a in soup.select("a.result__a")[:3]:
            title = a.get_text(strip=True)
            link = a["href"]
            # DuckDuckGo wraps real URL in a redirect, extract the real one
            m = re.search(r"uddg=(.+)", link)
            if m:
                link = httpx.URL(m.group(1)).decode()
            results.append({"title": title, "url": link})
        return json.dumps(results, ensure_ascii=False, indent=2)
    except Exception as e:
        return f"Error during web search: {e}"

@tool
def write_virtual_file(path: str, content: str) -> str:
    """
    Write a file into the virtual workspace. Path is relative to the workspace root.
    Returns a confirmation message.
    """
    try:
        full_path = WORKSPACE_DIR / path
        full_path.parent.mkdir(parents=True, exist_ok=True)
        full_path.write_text(content, encoding="utf-8")
        return f"File written to {path}"
    except Exception as e:
        return f"Failed to write file: {e}"

@tool
def list_virtual_files() -> str:
    """
    List all files currently present in the virtual workspace.
    Returns a newline separated list of relative paths.
    """
    try:
        files = [p.relative_to(WORKSPACE_DIR).as_posix() for p in WORKSPACE_DIR.rglob("*") if p.is_file()]
        return "\n".join(files) if files else "Workspace is empty."
    except Exception as e:
        return f"Error listing files: {e}"

@tool
def export_workspace() -> str:
    """
    Copy all files from the virtual workspace to the real output directory.
    Returns a summary of exported files.
    """
    try:
        if OUTPUT_DIR.exists():
            shutil.rmtree(OUTPUT_DIR)
        shutil.copytree(WORKSPACE_DIR, OUTPUT_DIR)
        exported = [p.relative_to(OUTPUT_DIR).as_posix() for p in OUTPUT_DIR.rglob("*") if p.is_file()]
        return "Exported files:\\n" + "\\n".join(exported)
    except Exception as e:
        return f"Export failed: {e}"

# ---------- Agent ----------
agent = create_deep_agent(
    model=LLM,
    tools=[web_search, write_virtual_file, list_virtual_files, export_workspace],
    backend=backend,
    system_prompt=(
        "You are a helpful research assistant. "
        "You can search the web, create virtual files, list them and finally export them to the real filesystem. "
        "When you have gathered enough information, write a summary to a file named 'report.txt' and then call export_workspace()."
    ),
)

# ---------- Main ----------
async def main():
    user_query = (
        "Find the latest information about the James Webb Space Telescope discoveries, "
        "summarize the top three findings, and save the summary to a file called 'jws_summary.txt' in the workspace. "
        "After that, export all virtual files to the real output directory."
    )
    result = await agent.ainvoke(
        {"messages": [HumanMessage(content=user_query)]},
        {"configurable": {"thread_id": "search-session-1"}},
    )
    # Print the final assistant message
    final_message = result["messages"][-1].content
    print("=== Final Assistant Message ===")
    print(final_message)

    # Show exported files
    if OUTPUT_DIR.exists():
        print("\\n=== Exported Files ===")
        for path in sorted(p.relative_to(OUTPUT_DIR).as_posix() for p in OUTPUT_DIR.rglob("*") if p.is_file()):
            print(path)

if __name__ == "__main__":
    asyncio.run(main())