import os
from typing import TypedDict, Dict, Any
import json

from langgraph import StateGraph, END
from langgraph.checkpoint.memory import MemorySaver
from langchain_openai import ChatOpenAI

# State definition
class ReflectState(TypedDict):
    question: str
    draft: str
    critique: str
    verdict: str  # "ok" | "needs_revision"
    round: int
    max_rounds: int

# LLM configuration

def get_llm() -> ChatOpenAI:
    """Return a ChatOpenAI instance configured via environment variables.
    If OLLAMA_BASE_URL is set, use Ollama; otherwise try OPENAI_API_KEY; fallback to local.
    """
    ollama_url = os.getenv("OLLAMA_BASE_URL")
    openai_key = os.getenv("OPENAI_API_KEY")
    if ollama_url:
        return ChatOpenAI(model="llama3", base_url=ollama_url, api_key="ollama")
    if openai_key:
        return ChatOpenAI(model="gpt-4o-mini")
    return ChatOpenAI(model="gpt-4o-mini")

llm = get_llm()

# Node: draft_answer

def draft_answer(state: ReflectState) -> Dict[str, Any]:
    question = state["question"]
    prompt = f"Answer the following question in 5-10 sentences: {question}"
    try:
        response = llm.invoke(prompt)
        draft = str(response)
    except Exception as e:
        raise RuntimeError(f"LLM draft generation failed: {e}")
    return {"draft": draft, "round": 1}

# Node: reflect

def reflect(state: ReflectState) -> Dict[str, Any]:
    draft = state["draft"]
    prompt = (
        "Critique the following draft for completeness, concreteness, and absence of filler. "
        "Respond with a JSON object containing 'verdict' (values: 'ok' or 'needs_revision') "
        "and 'critique' (text). Example: {\"verdict\": \"ok\", \"critique\": \"...\"}"
    )
    content = f"Draft: {draft}"
    full_prompt = prompt + "\n" + content
    try:
        response = llm.invoke(full_prompt)
        text = str(response)
    except Exception as e:
        raise RuntimeError(f"LLM reflect failed: {e}")
    # Parse JSON
    try:
        data = json.loads(text)
        verdict = data["verdict"]
        critique = data["critique"]
    except Exception:
        # Fallback parsing: treat entire response as critique
        verdict = "needs_revision"
        critique = text
    return {"critique": critique, "verdict": verdict}

# Node: rewrite

def rewrite(state: ReflectState) -> Dict[str, Any]:
    draft = state["draft"]
    critique = state["critique"]
    prompt = (
        "Rewrite the following draft based on the critique. "
        "Keep the meaning but improve clarity and remove filler. "
        "Output only the revised draft."
    )
    content = f"Draft: {draft}\nCritique: {critique}"
    full_prompt = prompt + "\n" + content
    try:
        response = llm.invoke(full_prompt)
        new_draft = str(response)
    except Exception as e:
        raise RuntimeError(f"LLM rewrite failed: {e}")
    return {"draft": new_draft, "round": state["round"] + 1}

# Build graph
builder = StateGraph(ReflectState)

builder.add_node("draft_answer", draft_answer)
builder.add_node("reflect", reflect)
builder.add_node("rewrite", rewrite)

builder.set_entry_point("draft_answer")

# Transitions
builder.add_edge("draft_answer", "reflect")
builder.add_conditional_edges(
    "reflect",
    lambda s: END if s["verdict"] == "ok" else "rewrite" if s["round"] < s["max_rounds"] else END,
)
builder.add_edge("rewrite", "reflect")

# Compile graph
graph = builder.compile(checkpointer=MemorySaver())

# CLI
if __name__ == "__main__":
    question = input("Enter your question: ")
    initial_state: ReflectState = {
        "question": question,
        "draft": "",
        "critique": "",
        "verdict": "",
        "round": 0,
        "max_rounds": 2,
    }
    result = graph.invoke(initial_state)
    print("\nFinal Result:\n")
    print(f"Draft: {result['draft']}")
    print(f"Critique: {result['critique']}")
    print(f"Verdict: {result['verdict']}")