"""LangGraph code review agent.

Implementation follows assignment:
- State: CodeReviewState with 4 criteria.
- Nodes: draft_review, reflect, rewrite.
- Graph: START -> draft_review -> reflect
  - ok -> END
  - needs_revision & round < max_rounds -> rewrite -> reflect
- Uses LangGraph and LangChain OpenAI for LLM calls.
- Structured output for critique via Pydantic model.
- Demo function sort_numbers.
"""

from __future__ import annotations

from typing import TypedDict, Dict

from langgraph.graph import StateGraph, END
from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage
from pydantic import BaseModel, Field

# ---------- State ----------
class CodeReviewState(TypedDict):
    code: str
    draft_review: str
    criteria_scores: Dict[str, int]
    weakest_criterion: str
    verdict: str
    round: int
    max_rounds: int

# ---------- LLM ----------
llm = ChatOpenAI(temperature=0)

# ---------- Nodes ----------
class DraftReviewOutput(BaseModel):
    review: str = Field(..., description="Draft review text")

class ReflectOutput(BaseModel):
    scores: Dict[str, int] = Field(..., description="Scores 0-10 for each criterion")
    weakest_criterion: str = Field(..., description="Criterion with lowest score")
    verdict: str = Field(..., description="'ok' or 'needs_revision'")


def draft_review(state: CodeReviewState) -> CodeReviewState:
    code = state["code"]
    prompt = f"""
    Write a concise code review (3-6 bullet points) for the following Python function. Focus on style, correctness, and potential improvements.

    ```python
    {code}
    ```
    """
    response = llm.invoke([HumanMessage(content=prompt)])
    data = DraftReviewOutput.model_validate_json(response.content)
    state["draft_review"] = data.review
    return state


def reflect(state: CodeReviewState) -> CodeReviewState:
    review = state["draft_review"]
    prompt = f"""
    Evaluate the following review on four criteria (PEP8, type_hints, edge_cases, naming) on a scale 0-10. Return a JSON object with keys: scores (dict), weakest_criterion (string), verdict ('ok' if all scores >=7 else 'needs_revision').

    Review:
    {review}
    """
    response = llm.invoke([HumanMessage(content=prompt)])
    data = ReflectOutput.model_validate_json(response.content)
    state["criteria_scores"] = data.scores
    state["weakest_criterion"] = data.weakest_criterion
    state["verdict"] = data.verdict
    return state


def rewrite(state: CodeReviewState) -> CodeReviewState:
    crit = state["weakest_criterion"]
    review = state["draft_review"]
    prompt = f"""
    The review below is weak in the {crit} criterion. Rewrite the review to improve that aspect.
    Original review:
    {review}
    """
    new_review = llm.invoke([HumanMessage(content=prompt)])
    state["draft_review"] = new_review.content
    state["round"] += 1
    return state

# ---------- Graph ----------
builder = StateGraph(CodeReviewState)
builder.add_node("draft_review", draft_review)
builder.add_node("reflect", reflect)
builder.add_node("rewrite", rewrite)

builder.set_entry_point("draft_review")
# After draft_review, go to reflect
builder.add_edge("draft_review", "reflect")
# After reflect, decide
builder.add_conditional_edges(
    "reflect",
    lambda state: state["verdict"] == "ok",
    {"ok": END, "needs_revision": "rewrite"},
)
# After rewrite, go back to reflect if rounds remain
builder.add_conditional_edges(
    "rewrite",
    lambda state: state["round"] < state["max_rounds"],
    {"continue": "reflect", "end": END},
)
# Compile graph

graph = builder.compile()

# ---------- Demo ----------

def sort_numbers(arr):
    return sorted(arr)

if __name__ == "__main__":
    code_str = "def sort_numbers(arr):\n    return sorted(arr)\n"
    initial_state: CodeReviewState = {
        "code": code_str,
        "draft_review": "",
        "criteria_scores": {},
        "weakest_criterion": "",
        "verdict": "",
        "round": 0,
        "max_rounds": 2,
    }
    result = graph.invoke(initial_state)
    print("Final state:")
    print(result)