"""LangGraph code review agent.

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

from __future__ import annotations

import os
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 ----------
# Use OpenAI only
llm = ChatOpenAI(temperature=0)

# ---------- Nodes ----------
class CritiqueOutput(BaseModel):
    review: str = Field(..., description="Draft review text")
    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 review_and_critique(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}
    ```

    Then evaluate the review on four criteria (PEP8, type_hints, edge_cases, naming) on a scale 0-10.
    Return a JSON object with keys: review (string), scores (dict), weakest_criterion (string), verdict ('ok' if all scores >=7 else 'needs_revision').
    """
    response = llm.invoke([HumanMessage(content=prompt)])
    data = CritiqueOutput.model_validate_json(response.content)
    state["draft_review"] = data.review
    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("review_and_critique", review_and_critique)
builder.add_node("rewrite", rewrite)

builder.set_entry_point("review_and_critique")
# After initial review_and_critique, decide to end if verdict ok
builder.add_conditional_edges(
    "review_and_critique",
    lambda state: state["verdict"] == "ok",
    {"ok": END, "needs_revision": "rewrite"},
)
# After rewrite, go back to review_and_critique if rounds remain
builder.add_conditional_edges(
    "rewrite",
    lambda state: "review_and_critique" if state["round"] < state["max_rounds"] else END,
    {"review_and_critique": "review_and_critique", END: END}
)

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)