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

from dotenv import load_dotenv
from langchain_openai import ChatOpenAI, OpenAIEmbeddings
from langchain_core.messages import HumanMessage, SystemMessage
from langchain_core.prompts import ChatPromptTemplate
from langchain.tools import tool
from tavily import TavilyClient

from deepagents import create_deep_agent
from deepagents.backends import CompositeBackend, LocalShellBackend, FilesystemBackend

from langgraph.graph import StateGraph, START, END
from langgraph.graph.message import add_messages

load_dotenv()

# -------------------- LLM --------------------
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 tool --------------------
tavily = TavilyClient(api_key=os.getenv("TAVILY_API_KEY"))


@tool
def web_search(query: str) -> str:
    """
    Perform a web search using Tavily and return a concise summary of the top result.
    """
    try:
        response = tavily.search(query, search_depth="basic", max_results=3)
        results = response.get("results", [])
        if not results:
            return "No relevant results found."
        # Concatenate titles and snippets
        summary = " ".join(r.get("title", "") + ". " + r.get("content", "") for r in results[:2])
        return summary.strip()
    except Exception as e:
        return f"Search error: {e}"


# -------------------- DeepAgent wrapper (required by course) --------------------
backend = CompositeBackend(
    [
        LocalShellBackend(workspace_dir="./workspace"),
        FilesystemBackend(),
    ]
)

deep_agent = create_deep_agent(
    model=llm,
    tools=[web_search],
    backend=backend,
    system_prompt="You are a helpful research assistant.",
)


async def invoke_llm(messages: List[Dict[str, str]]) -> str:
    """
    Helper that sends messages to the deep agent and returns the assistant's reply.
    """
    result = await deep_agent.ainvoke(
        {"messages": [HumanMessage(content=m["content"]) for m in messages]},
        {"configurable": {"thread_id": "compare-session"}},
    )
    return result["messages"][-1].content


# -------------------- State definition --------------------
class CompareState(TypedDict):
    entities: List[str]                     # 3 names to compare
    criteria: List[str]                     # 3-5 criteria
    pairs: List[Tuple[str, str]]            # remaining (entity, criterion) pairs
    findings: Dict[str, List[str]]          # entity -> list of notes (same order as criteria)
    final_table: str | None
    verdict: str | None
    messages: Annotated[list, add_messages]  # for LangGraph internal use


# -------------------- Node: plan criteria --------------------
async def plan_criteria(state: CompareState) -> CompareState:
    prompt = ChatPromptTemplate.from_messages(
        [
            SystemMessage(
                content="You are an expert analyst. Given three entities, propose 3 to 5 criteria to compare them. Return the criteria as a JSON list."
            ),
            HumanMessage(content=f"Entities: {', '.join(state['entities'])}"),
        ]
    )
    response = await invoke_llm([{"role": "system", "content": prompt.messages[0].content},
                                {"role": "user", "content": prompt.messages[1].content}])
    try:
        import json

        criteria = json.loads(response)
        if not isinstance(criteria, list):
            raise ValueError
    except Exception:
        # Fallback: split by newlines
        criteria = [c.strip("- ").strip() for c in response.splitlines() if c.strip()]
    state["criteria"] = criteria[:5]
    # Build all pairs
    state["pairs"] = [(e, c) for e in state["entities"] for c in state["criteria"]]
    # Initialise findings dict
    state["findings"] = {e: [] for e in state["entities"]}
    return state


# -------------------- Node: research entity --------------------
async def research_entity(state: CompareState) -> CompareState:
    if not state["pairs"]:
        return state
    entity, criterion = state["pairs"].pop(0)
    query = f"{entity} {criterion}"
    # Use the web_search tool directly (synchronous call is fine)
    note = web_search.run(query)  # type: ignore
    # Append note to the correct entity list
    state["findings"][entity].append(note)
    # Log for CLI
    print(f"[{entity} × {criterion}] найдено: {note[:200]}...")
    return state


# -------------------- Node: build table --------------------
def build_table(state: CompareState) -> CompareState:
    header = ["Критерий"] + state["entities"]
    rows = []
    for idx, crit in enumerate(state["criteria"]):
        row = [crit]
        for ent in state["entities"]:
            notes = state["findings"].get(ent, [])
            note = notes[idx] if idx < len(notes) else ""
            row.append(note.replace("\n", " ").strip())
        rows.append(row)

    # Markdown table construction
    def md_row(cols: List[str]) -> str:
        return "| " + " | ".join(cols) + " |"

    separator = "| " + " | ".join(["---"] * len(header)) + " |"
    table_lines = [md_row(header), separator] + [md_row(r) for r in rows]
    state["final_table"] = "\n".join(table_lines)
    return state


# -------------------- Node: verdict --------------------
async def verdict(state: CompareState) -> CompareState:
    prompt = ChatPromptTemplate.from_messages(
        [
            SystemMessage(
                content="You are an analyst. Based on the comparison table, give a short verdict (2-4 sentences) recommending which entity is best for which use case."
            ),
            HumanMessage(content=state["final_table"] or ""),
        ]
    )
    response = await invoke_llm([{"role": "system", "content": prompt.messages[0].content},
                                {"role": "user", "content": prompt.messages[1].content}])
    state["verdict"] = response.strip()
    return state


# -------------------- Graph assembly --------------------
workflow = StateGraph(CompareState)

workflow.add_node("plan_criteria", plan_criteria)
workflow.add_node("research_entity", research_entity)
workflow.add_node("build_table", build_table)
workflow.add_node("verdict", verdict)

workflow.add_edge(START, "plan_criteria")
workflow.add_conditional_edges(
    "plan_criteria",
    lambda s: "research_entity" if s["pairs"] else "build_table",
)

workflow.add_edge("research_entity", "research_entity")  # loop until pairs empty
workflow.add_conditional_edges(
    "research_entity",
    lambda s: "research_entity" if s["pairs"] else "build_table",
)

workflow.add_edge("build_table", "verdict")
workflow.add_edge("verdict", END)

app = workflow.compile()


# -------------------- CLI demo --------------------
async def main() -> None:
    default_entities = ["Chroma", "FAISS", "Qdrant"]
    user_input = input(
        "Enter three entities separated by commas (or press Enter for default): "
    ).strip()
    entities = (
        [e.strip() for e in user_input.split(",") if e.strip()]
        if user_input
        else default_entities
    )
    if len(entities) != 3:
        print("Please provide exactly three entities.")
        return

    initial_state: CompareState = {
        "entities": entities,
        "criteria": [],
        "pairs": [],
        "findings": {},
        "final_table": None,
        "verdict": None,
        "messages": [],
    }

    # Run the graph
    async for event in app.astream(initial_state):
        # The graph updates state internally; we only need final output after END
        pass

    final = event  # last state after END
    print("\n=== План критериев ===")
    print(", ".join(final["criteria"]))
    print("\n=== Итоговая таблица ===")
    print(final["final_table"])
    print("\n=== Вердикт ===")
    print(final["verdict"])


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