import os
import argparse
import asyncio
from typing import TypedDict, List, Dict, Optional

from dotenv import load_dotenv
from langchain_openai import ChatOpenAI
from langchain.tools import tool
from deepagents import create_deep_agent
from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend

from langgraph.graph import StateGraph, START, END
from langchain_core.output_parsers import PydanticOutputParser
from pydantic import BaseModel, Field

from tavily import TavilyClient

# Load environment variables
load_dotenv()
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
TAVILY_API_KEY = os.getenv("TAVILY_API_KEY")

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

# Tavily client
tavily_client = TavilyClient(api_key=TAVILY_API_KEY)

# State definition
class CompareState(TypedDict):
    entities: List[str]
    criteria: List[str]
    findings: Dict[str, List[str]]
    final_table: Optional[str]
    verdict: Optional[str]

# Pydantic models for parsing
class CriteriaOutput(BaseModel):
    criteria: List[str] = Field(description="List of comparison criteria")

class VerdictOutput(BaseModel):
    verdict: str = Field(description="Verdict text")

criteria_parser = PydanticOutputParser(pydantic_object=CriteriaOutput)
verdict_parser = PydanticOutputParser(pydantic_object=VerdictOutput)

# Node: plan_criteria
def plan_criteria(state: CompareState) -> CompareState:
    prompt = (
        f"Given the following entities: {', '.join(state['entities'])}. "
        "Generate 3 to 5 distinct criteria for comparing these entities. "
        "Return a JSON object with a field 'criteria' that is a list of strings."
    )
    response = llm.invoke(prompt)
    parsed = criteria_parser.parse(response.content)
    state["criteria"] = parsed.criteria
    # Initialize findings dict
    state["findings"] = {entity: [] for entity in state["entities"]}
    return state

# Node: research_entity
def research_entity(state: CompareState) -> CompareState:
    # Find first entity with missing notes
    for entity in state["entities"]:
        if len(state["findings"][entity]) < len(state["criteria"]):
            idx = len(state["findings"][entity])
            criterion = state["criteria"][idx]
            query = f"{entity} {criterion}"
            # Perform Tavily search
            results = tavily_client.search(query)
            # Take first result snippet
            if results and results[0].snippet:
                note = results[0].snippet.strip()
            else:
                note = f"No relevant info found for {entity} on {criterion}."
            state["findings"][entity].append(note)
            break
    return state

# Node: check_done
def check_done(state: CompareState) -> str:
    for entity in state["entities"]:
        if len(state["findings"][entity]) < len(state["criteria"]):
            return "continue"
    return "done"

# Node: build_table
def build_table(state: CompareState) -> CompareState:
    header = ["Criterion"] + state["entities"]
    rows = []
    for criterion in state["criteria"]:
        row = [criterion]
        for entity in state["entities"]:
            notes = state["findings"][entity]
            idx = state["criteria"].index(criterion)
            note = notes[idx] if idx < len(notes) else ""
            row.append(note)
        rows.append(row)
    # Build markdown table
    table_lines = ["| " + " | ".join(header) + " |"]
    table_lines.append("|" + "|".join(["---"] * len(header)) + "|")
    for row in rows:
        table_lines.append("| " + " | ".join(row) + " |")
    table_md = "\n".join(table_lines)
    state["final_table"] = table_md
    return state

# Node: verdict
def verdict(state: CompareState) -> CompareState:
    prompt = (
        f"Here is a markdown table comparing the entities:\n\n{state['final_table']}\n\n"
        "Based on this table, provide a concise verdict recommending which entity "
        "is best suited for a typical vector database use case. "
        "Return a JSON object with a field 'verdict' that is a string."
    )
    response = llm.invoke(prompt)
    parsed = verdict_parser.parse(response.content)
    state["verdict"] = parsed.verdict
    return state

# Build LangGraph
graph = StateGraph(CompareState)
graph.add_node("plan_criteria", plan_criteria)
graph.add_node("research_entity", research_entity)
graph.add_node("check_done", check_done)
graph.add_node("build_table", build_table)
graph.add_node("verdict", verdict)

graph.add_edge(START, "plan_criteria")
graph.add_edge("plan_criteria", "research_entity")
graph.add_edge("research_entity", "check_done")
graph.add_conditional_edges(
    "check_done",
    lambda x: x,
    {
        "continue": "research_entity",
        "done": "build_table",
    },
)
graph.add_edge("build_table", "verdict")
graph.add_edge("verdict", END)

compiled_graph = graph.compile()

# Tool: compare_entities
@tool
def compare_entities(query: str) -> str:
    """
    Compare three entities based on user query.
    Expected format: "Compare 3 vector DBs: Chroma, FAISS, Qdrant"
    """
    # Extract entities after colon
    if ":" in query:
        parts = query.split(":", 1)
        entities_part = parts[1]
    else:
        entities_part = query
    entities = [e.strip() for e in entities_part.split(",") if e.strip()]
    if len(entities) != 3:
        return "Please provide exactly three entities separated by commas."
    initial_state: CompareState = {
        "entities": entities,
        "criteria": [],
        "findings": {},
        "final_table": None,
        "verdict": None,
    }
    final_state = compiled_graph.invoke(initial_state)
    table = final_state["final_table"] or ""
    verdict_text = final_state["verdict"] or ""
    return f"{table}\n\nVerdict:\n{verdict_text}"

# DeepAgent setup
backend = CompositeBackend([LocalShellBackend(workspace_dir="./workspace"), FilesystemBackend()])

agent = create_deep_agent(
    model=llm,
    tools=[compare_entities],
    backend=backend,
    system_prompt="You are a helpful assistant that can compare three entities.",
)

# CLI
def main():
    parser = argparse.ArgumentParser(description="Compare three entities.")
    parser.add_argument(
        "--query",
        type=str,
        default="Compare 3 vector DBs: Chroma, FAISS, Qdrant",
        help="Comparison query in the format 'Compare 3 vector DBs: A, B, C'",
    )
    args = parser.parse_args()
    async def run():
        result = await agent.ainvoke(
            {"messages": [{"role": "user", "content": args.query}]},
            {"configurable": {"thread_id": "session-1"}},
        )
        print(result["messages"][-1]["content"])
    asyncio.run(run())

if __name__ == "__main__":
    main()