import os
from typing import Dict, List, Tuple

from langgraph.graph import StateGraph, END
from langchain_openai import ChatOpenAI
from tavily import TavilyClient
from dotenv import load_dotenv

from .state import CompareState

load_dotenv()

# Initialize LLM and Tavily client
llm = ChatOpenAI(
    temperature=0.2,
    model="gpt-4o-mini",
    openai_api_key=os.getenv("OPENAI_API_KEY"),
)

tavily = TavilyClient(api_key=os.getenv("TAVILY_API_KEY"))

def plan_criteria(state: CompareState) -> CompareState:
    """
    Generate 3–5 comparison criteria for the given entities.
    """
    entities = state.get("entities", [])
    if not entities:
        raise ValueError("No entities provided for criteria planning.")

    prompt = (
        f"Given the following entities: {', '.join(entities)}.\n"
        "Suggest 3 to 5 key criteria to compare them. "
        "Return the criteria as a numbered list, one per line."
    )
    response = llm.invoke(prompt)
    criteria_text = response.content.strip()
    # Parse numbered list
    criteria = []
    for line in criteria_text.splitlines():
        line = line.strip()
        if line:
            # Remove leading numbers if present
            if line[0].isdigit() and (len(line) > 1 and line[1] in ". "):
                line = line[2:].strip()
            criteria.append(line)
    state["criteria"] = criteria
    return state

def research_entity(state: CompareState) -> CompareState:
    """
    For each entity–criterion pair, perform a Tavily web search
    and store a short note in findings.
    """
    entities = state.get("entities", [])
    criteria = state.get("criteria", [])
    findings: Dict[str, List[str]] = {entity: [] for entity in entities}

    for entity in entities:
        for criterion in criteria:
            query = f"{entity} {criterion}"
            try:
                result = tavily.search(query=query, max_results=1)
                if result and result["results"]:
                    snippet = result["results"][0]["content"][:200]
                else:
                    snippet = "No relevant information found."
            except Exception as e:
                snippet = f"Error during search: {e}"
            findings[entity].append(snippet)

    state["findings"] = findings
    return state

def build_table(state: CompareState) -> CompareState:
    """
    Build a Markdown table from findings.
    Rows: criteria, Columns: entities.
    """
    entities = state.get("entities", [])
    criteria = state.get("criteria", [])
    findings = state.get("findings", {})

    header = "| Criterion | " + " | ".join(entities) + " |\n"
    separator = "|---" * (len(entities) + 1) + "|\n"

    rows = ""
    for idx, criterion in enumerate(criteria):
        row = f"| {criterion} | "
        for entity in entities:
            notes = findings.get(entity, [])
            note = notes[idx] if idx < len(notes) else ""
            # Escape pipe characters
            note = note.replace("|", "\\|")
            row += f"{note} | "
        rows += row + "\n"

    table = header + separator + rows
    state["final_table"] = table
    return state

def verdict(state: CompareState) -> CompareState:
    """
    Generate a verdict recommendation based on the table and criteria.
    """
    table = state.get("final_table", "")
    criteria = state.get("criteria", [])
    entities = state.get("entities", [])

    prompt = (
        f"Here is a comparative table of the following entities: {', '.join(entities)}.\n\n"
        f"{table}\n\n"
        f"Based on the criteria: {', '.join(criteria)}.\n"
        "Provide a concise recommendation (2–4 sentences) indicating which entity is best suited for which use case."
    )
    response = llm.invoke(prompt)
    state["verdict"] = response.content.strip()
    return state