import os
from dotenv import load_dotenv
from typing import List
from langchain_community.tools import DuckDuckGoSearchRun
from langchain_experimental.tools import PythonREPLTool
from langchain_community.vectorstores import FAISS
#from langchain.embeddings import OpenAIEmbeddings
from langchain_core.documents import Document
from langchain_huggingface import HuggingFaceEndpoint, ChatHuggingFace, HuggingFaceEmbeddings
from langchain_core.messages import HumanMessage



load_dotenv()

# -----------------------------
# LLM
# -----------------------------
#llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)
#repo_id="deepseek-ai/DeepSeek-V4-Pro"
llm = ChatHuggingFace(
    llm=HuggingFaceEndpoint(
        repo_id="Qwen/Qwen2.5-Coder-32B-Instruct",
        huggingfacehub_api_token='HF_KEY',
        task="conversational", # Specify task for the conversational model
    )
)

# -----------------------------
# Tools
# -----------------------------
search = DuckDuckGoSearchRun()
python_tool = PythonREPLTool()

TOOLS = {
    "search": search.run,
    "python": python_tool.run,
    "llm": lambda x: llm.invoke([HumanMessage(content=x)]).content,
    "summarize": lambda text: llm.invoke([HumanMessage(content=f"Summarize the following:\n{text}")]).content
}

# -----------------------------
# Memory (Vector DB)
# -----------------------------
embeddings = HuggingFaceEmbeddings()
# Initialize FAISS with a dummy document to prevent IndexError when trying to determine embedding dimension
vectorstore = FAISS.from_documents([Document(page_content="initialization_document_for_dimension_inference")], embeddings)

def store_memory(text: str):
    vectorstore.add_documents([Document(page_content=text)])

def retrieve_memory(query: str):
    docs = vectorstore.similarity_search(query, k=3)
    return "\n".join([d.page_content for d in docs])

# -----------------------------
# Planner
# -----------------------------
def plan(goal, history):
    prompt = f"""
You are an autonomous agent.

Goal: {goal}

Previous steps:
{history}

Decide the NEXT action:
- search(query)
- python(code)
- llm(prompt)
- summarize(text)
- finish(answer)

Respond ONLY in one line.
"""
    return llm.invoke([HumanMessage(content=prompt)]).content.strip()

# -----------------------------
# Executor
# -----------------------------
def execute(action: str):
    try:
        if action.startswith("search("):
            query = action[len("search("):-1]
            return TOOLS["search"](query)

        elif action.startswith("python("):
            code = action[len("python("):-1]
            return TOOLS["python"](code)

        elif action.startswith("llm("):
            prompt = action[len("llm("):-1]
            return TOOLS["llm"](prompt)

        elif action.startswith("summarize("):
            text_to_summarize = action[len("summarize("):-1]
            return TOOLS["summarize"](text_to_summarize)

        elif action.startswith("finish("):
            return action[len("finish("):-1]

        else:
            return "Invalid action"

    except Exception as e:
        return f"Error: {str(e)}"

# -----------------------------
# Critic (loop control)
# -----------------------------
def critic(goal, last_result):
    prompt = f"""
Goal: {goal}

Latest result:
{last_result}

Is the goal achieved? Answer YES or NO.
"""
    return "YES" in llm.invoke([HumanMessage(content=prompt)]).content.upper()

# -----------------------------
# Autonomous Loop
# -----------------------------
def autonomous_agent(goal: str, max_steps=15):

    history = ""
    print(f"\n🎯 Goal: {goal}\n")

    for step in range(max_steps):
        print(f"--- Step {step+1} ---")

        # Retrieve memory
        memory_context = retrieve_memory(goal)

        action = plan(goal, history + "\nMemory:\n" + memory_context)
        print(f"🧠 Plan: {action}")

        result = execute(action)
        print(f"⚙️ Result: {result[:300]}...\n")

        # Store memory
        store_memory(f"Action: {action}\nResult: {result}")

        history += f"\nStep {step+1}: {action} → {result}"

        # Finish condition
        if action.startswith("finish("):
            print("✅ Finished by agent")
            return result

        # Critic check
        if critic(goal, result):
            print("✅ Critic determined goal achieved")
            return result

    return "❌ Max steps reached without completion"

# -----------------------------
# Run
# -----------------------------
if __name__ == "__main__":
    while True:
        goal = input("\nEnter goal (or 'exit'): ")
        if goal == "exit":
            break

        result = autonomous_agent(goal)
        print(f"\n🤖 Final Output:\n{result}\n")