import os
from typing import Dict, List
from dotenv import load_dotenv
from langchain_openai import ChatOpenAI
from langchain_tavily import TavilySearchTool
from langgraph.prebuilt import create_chat_agent
from langgraph import add_messages, StateGraph
from src.state import CompareState

load_dotenv()

# LLM and Tavily tool
llm = ChatOpenAI(temperature=0.7)
tavily = TavilySearchTool(api_key=os.getenv("TAVILY_API_KEY"))

# Helper to format findings for LLM
def format_findings(findings: Dict[str, List[str]]) -> str:
    parts = []
    for entity, notes in findings.items():
        parts.append(f"**{entity}**:")
        for note in notes:
            parts.append(f"- {note}")
    return "\n".join(parts)

# Node: Generate comparison criteria
def plan_criteria(state: CompareState) -> CompareState:
    entities = state["entities"]
    prompt = (
        f"Generate 3-5 concise comparison criteria for the following entities: "
        f"{', '.join(entities)}. Return a numbered list."
    )
    response = llm.invoke(prompt)
    # Parse numbered list
    criteria = []
    for line in response.splitlines():
        line = line.strip()
        if line:
            # Remove leading numbers
            if line[0].isdigit() and (len(line) > 1 and line[1] in ". "):
                line = line[2:].strip()
            criteria.append(line)
    state["criteria"] = criteria
    # Initialize findings dict
    state["findings"] = {entity: [] for entity in entities}
    return state

# Node: Research one entity-criterion pair
def research_entity(state: CompareState) -> CompareState:
    entities = state["entities"]
    criteria = state["criteria"]
    findings = state["findings"]

    # Find next entity needing research
    for entity in entities:
        if len(findings[entity]) < len(criteria):
            idx = len(findings[entity])
            criterion = criteria[idx]
            query = f"{entity} {criterion}"
            # Tavily search
            results = tavily.invoke({"query": query, "max_results": 3})
            # Take first result snippet
            if results and "results" in results and len(results["results"]) > 0:
                snippet = results["results"][0]["snippet"]
                source = results["results"][0]["url"]
                note = f"{criterion}: {snippet} (Source: {source})"
            else:
                note = f"{criterion}: No recent information found."
            findings[entity].append(note)
            break
    state["findings"] = findings
    return state

# Node: Build cohesive research brief
def build_brief(state: CompareState) -> CompareState:
    findings_text = format_findings(state["findings"])
    prompt = (
        f"Using the following findings, write a cohesive research brief that summarizes "
        f"the strengths and weaknesses of each entity. The brief should be clear, "
        f"structured, and suitable for a technical audience.\n\n"
        f"Findings:\n{findings_text}"
    )
    brief = llm.invoke(prompt)
    state["final_brief"] = brief
    return state

# Node: Generate verdict/recommendation
def verdict(state: CompareState) -> CompareState:
    brief = state["final_brief"]
    criteria = state["criteria"]
    prompt = (
        f"Based on the research brief below and the comparison criteria, provide a "
        f"clear recommendation on which entity is best suited for a typical use case. "
        f"Explain your reasoning in 2-4 sentences.\n\n"
        f"Research Brief:\n{brief}\n\n"
        f"Criteria:\n- " + "\n- ".join(criteria)
    )
    recommendation = llm.invoke(prompt)
    state["verdict"] = recommendation
    return state