"""Deep Agent implementation based on LangGraph.

This agent can:
1. Search the web for information using DuckDuckGo API.
2. Create virtual files in memory.
3. At the end of the run, persist virtual files to the real file system.

The code follows the requirements:
- Uses langchain>=1.0.0 and langgraph>=1.0.0.
- Correct imports for text splitters, Chroma, Ollama embeddings and chat model.
- Implements `create_agent` and `create_agent_executor` functions.
- Uses the `write_file` tool to write virtual files.
"""

import pathlib
from typing import Dict, List, Any, Optional

# LangChain imports
from langchain_ollama import ChatOllama, OllamaEmbeddings
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_chroma import Chroma

# LangGraph imports
from langgraph.graph import Graph, State

# Requests for web search
import requests

# -----------------------------
# Configuration
# -----------------------------
OLLAMA_MODEL = "llama3"
EMBEDDINGS_MODEL = "llama3"
SEARCH_URL = "https://api.duckduckgo.com/"
OUTPUT_DIR = pathlib.Path("output_files")
OUTPUT_DIR.mkdir(exist_ok=True)

# -----------------------------
# Helper functions
# -----------------------------

def search_web(query: str) -> str:
    """Return a short summary of the search results using DuckDuckGo."""
    params = {
        "q": query,
        "format": "json",
        "no_html": 1,
        "skip_disambig": 1,
    }
    try:
        r = requests.get(SEARCH_URL, params=params, timeout=10)
        r.raise_for_status()
        data = r.json()
        abstract = data.get("AbstractText")
        if abstract:
            return abstract
        topics = data.get("RelatedTopics", [])
        snippets = [t.get("Text", "") for t in topics if "Text" in t]
        return "\n".join(snippets[:5])
    except Exception as e:
        return f"Error during search: {e}"

# -----------------------------
# State definition
# -----------------------------
class AgentState(State):
    query: str
    virtual_files: Dict[str, str]
    history: List[Dict[str, str]]
    answer: Optional[str] = None

# -----------------------------
# Tool: write_file
# -----------------------------

def write_file_tool(state: AgentState, file_name: str, content: str) -> AgentState:
    new_files = state.virtual_files.copy()
    new_files[file_name] = content
    return state.copy(update={"virtual_files": new_files})

# -----------------------------
# Tool: search
# -----------------------------

def search_tool(state: AgentState, query: str) -> AgentState:
    result = search_web(query)
    new_history = state.history + [{"role": "tool", "name": "search", "content": result}]
    return state.copy(update={"history": new_history})

# -----------------------------
# Agent logic
# -----------------------------

def create_agent() -> Graph:
    llm = ChatOllama(model=OLLAMA_MODEL, temperature=0.7)
    embeddings = OllamaEmbeddings(model=EMBEDDINGS_MODEL)
    splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)

    def agent(state: AgentState) -> Dict[str, Any]:
        # Build prompt from history
        messages = []
        for msg in state.history:
            if msg["role"] == "user":
                messages.append({"role": "user", "content": msg["content"]})
            elif msg["role"] == "assistant":
                messages.append({"role": "assistant", "content": msg["content"]})
            elif msg["role"] == "tool":
                messages.append({"role": "assistant", "content": f"[Tool: {msg['name']}] {msg['content']}"})
        messages.append({"role": "assistant", "content": f"User query: {state.query}"})

        response = llm.invoke(messages)
        text = response.content.strip()

        if text.upper().startswith("ANSWER:"):
            answer = text[7:].strip()
            new_history = state.history + [{"role": "assistant", "content": answer}]
            return {"final_answer": answer, "history": new_history}
        elif text.upper().startswith("SEARCH:"):
            query = text[7:].strip()
            return {"search_query": query}
        elif text.upper().startswith("WRITE:"):
            try:
                rest = text[6:].strip()
                file_name, content = rest.split("|", 1)
                return {"write_file": {"file_name": file_name.strip(), "content": content.strip()}}
            except Exception:
                return {"final_answer": "Could not parse WRITE command."}
        else:
            answer = text
            new_history = state.history + [{"role": "assistant", "content": answer}]
            return {"final_answer": answer, "history": new_history}

    graph = Graph()
    graph.add_node("agent", agent)
    graph.add_node("search", search_tool)
    graph.add_node("write_file", write_file_tool)

    def final_answer_node(state: AgentState) -> AgentState:
        for msg in reversed(state.history):
            if msg["role"] == "assistant":
                state = state.copy(update={"answer": msg["content"]})
                break
        return state

    graph.add_node("final_answer", final_answer_node)

    graph.add_edge("agent", "search", condition=lambda out: "search_query" in out)
    graph.add_edge("agent", "write_file", condition=lambda out: "write_file" in out)
    graph.add_edge("agent", "final_answer", condition=lambda out: "final_answer" in out)
    graph.add_edge("search", "agent", condition=lambda out: True)
    graph.add_edge("write_file", "agent", condition=lambda out: True)

    graph.set_start("agent")
    graph.set_end("final_answer")

    return graph

# -----------------------------
# Executor helper
# -----------------------------

def create_agent_executor() -> Graph:
    return create_agent()

# -----------------------------
# Main execution
# -----------------------------
if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser(description="Run the deep agent.")
    parser.add_argument("query", type=str, help="User query to process")
    args = parser.parse_args()

    graph = create_agent_executor()
    initial_state = AgentState(query=args.query, virtual_files={}, history=[{"role": "user", "content": args.query}])

    final_state = graph.invoke(initial_state)

    print("\n=== Final Answer ===\n")
    print(final_state.answer if final_state.answer else "No answer produced.")

    for fname, content in final_state.virtual_files.items():
        out_path = OUTPUT_DIR / fname
        out_path.parent.mkdir(parents=True, exist_ok=True)
        with open(out_path, "w", encoding="utf-8") as f:
            f.write(content)
        print(f"Virtual file written to {out_path}")