import os
import asyncio
import argparse
import re
from typing import TypedDict, Annotated, Dict, List, Tuple, Any

from dotenv import load_dotenv
from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage
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 tavily import TavilySearchResults

# Load environment variables
load_dotenv()

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

# Tavily client
tavily = TavilySearchResults(api_key=os.getenv("TAVILY_API_KEY"))

# State definition
class CompareState(TypedDict):
    entities: List[str]                # 3 names for comparison
    criteria: List[str]                # 3-5 criteria
    findings: Dict[str, Dict[str, str]]  # entity -> criterion -> note
    pairs_to_process: List[Tuple[str, str]]  # (entity, criterion)
    final_table: str | None
    verdict: str | None

# Node: plan_criteria
def plan_criteria(state: CompareState) -> CompareState:
    prompt = (
        f"Given the entities {state['entities']}, generate 3 to 5 comparison criteria. "
        "Return a JSON array of strings."
    )
    response = llm.invoke(prompt)
    # Extract JSON array
    try:
        import json
        criteria = json.loads(response.content)
        if not isinstance(criteria, list):
            raise ValueError
    except Exception:
        criteria = ["performance", "scalability", "ease of use"]
    state["criteria"] = criteria
    # Prepare pairs to process
    state["pairs_to_process"] = [(entity, criterion) for entity in state["entities"] for criterion in criteria]
    state["findings"] = {entity: {} for entity in state["entities"]}
    print(f"[plan_criteria] Generated criteria: {criteria}")
    return state

# Node: research_entity
def research_entity(state: CompareState) -> CompareState:
    if not state["pairs_to_process"]:
        return state
    entity, criterion = state["pairs_to_process"].pop(0)
    query = f"{entity} {criterion}"
    results = tavily.search(query)
    snippet = results[0].content if results else "No relevant information found."
    note = f"{criterion}: {snippet}"
    state["findings"][entity][criterion] = note
    print(f"[research_entity] ({entity} × {criterion}) found: {snippet[:60]}...")
    return state

# Node: check_pairs
def check_pairs(state: CompareState) -> str:
    return "continue" if state["pairs_to_process"] else "done"

# Node: build_table
def build_table(state: CompareState) -> CompareState:
    headers = ["Criterion"] + state["entities"]
    table_rows = []
    for criterion in state["criteria"]:
        row = [criterion]
        for entity in state["entities"]:
            note = state["findings"][entity].get(criterion, "N/A")
            row.append(note)
        table_rows.append(row)
    # Build markdown table
    md = "| " + " | ".join(headers) + " |\n"
    md += "| " + " | ".join(["---"] * len(headers)) + " |\n"
    for row in table_rows:
        md += "| " + " | ".join(row) + " |\n"
    state["final_table"] = md
    print("[build_table] Table constructed.")
    return state

# Node: verdict
def verdict(state: CompareState) -> CompareState:
    prompt = (
        f"Based on the following table, provide a concise verdict recommending which entity is best for which use case:\n\n"
        f"{state['final_table']}\n\n"
        "Answer in 2-4 sentences."
    )
    response = llm.invoke(prompt)
    state["verdict"] = response.content.strip()
    print("[verdict] Verdict generated.")
    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_pairs", check_pairs)
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_pairs")
graph.add_conditional_edges(
    "check_pairs",
    lambda x: x,
    {
        "continue": "research_entity",
        "done": "build_table",
    },
)
graph.add_edge("build_table", "verdict")
graph.add_edge("verdict", END)

app = graph.compile()

# Tool that runs the comparison
@tool
def run_comparison(query: str) -> str:
    """
    Run a comparative review of three entities.
    Input: a string containing three entity names separated by commas.
    Output: markdown table and verdict.
    """
    # Extract entities
    entities = [e.strip() for e in re.split(r",|;|and", query) if e.strip()]
    if len(entities) != 3:
        return "Please provide exactly three entities separated by commas."
    initial_state: CompareState = {
        "entities": entities,
        "criteria": [],
        "findings": {},
        "pairs_to_process": [],
        "final_table": None,
        "verdict": None,
    }
    result_state = app.invoke(initial_state)
    table = result_state["final_table"] or ""
    verdict_text = result_state["verdict"] or ""
    return f"{table}\n\n**Verdict:**\n{verdict_text}"

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

# Create deepagents agent
agent = create_deep_agent(
    model=llm,
    tools=[run_comparison],
    backend=backend,
    system_prompt="You are a helpful agent that performs comparative reviews.",
)

async def main():
    parser = argparse.ArgumentParser(description="Comparative review CLI")
    parser.add_argument(
        "-e",
        "--entities",
        type=str,
        help="Comma-separated list of three entities to compare",
    )
    args = parser.parse_args()

    if args.entities:
        query = args.entities
    else:
        # Default entities
        query = "Chroma, FAISS, Qdrant"

    print(f"Running comparison for: {query}")

    result = await agent.ainvoke(
        {"messages": [HumanMessage(content=query)]},
        {"configurable": {"thread_id": "session-1"}},
    )
    final_output = result["messages"][-1].content
    print("\n=== Final Output ===\n")
    print(final_output)

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