"""LangGraph Code Review Agent with Reflection.

This repository implements a LangGraph agent that takes a Python function as input and
produces a code review. A critic node evaluates the review on four criteria:

* PEP8 compliance
* Type hints
* Edge cases handling
* Naming conventions

If the critic returns `needs_revision`, the `rewrite` node rewrites the weakest
criterion section of the review. The process repeats until the critic is satisfied
or the maximum number of rounds is reached.

The demo in ``__main__`` shows how to run the agent on a simple function.
"""

from __future__ import annotations

import os
from typing import TypedDict, Dict

from langchain_openai import ChatOpenAI
from langgraph.graph import StateGraph, END
from langgraph.prebuilt import create_react_agent
from pydantic import BaseModel, Field

# ---------------------------------------------------------------------------
# State definition
# ---------------------------------------------------------------------------
class CodeReviewState(TypedDict):
    code: str
    draft_review: str | None
    criteria_scores: Dict[str, int] | None
    weakest_criterion: str | None
    verdict: str | None  # "ok" | "needs_revision"
    round: int
    max_rounds: int

# ---------------------------------------------------------------------------
# LLM configuration
# ---------------------------------------------------------------------------
# The OpenAI API key must be set in the environment variable OPENAI_API_KEY.
# For local Ollama usage, replace ChatOpenAI with ChatOllama.
llm = ChatOpenAI(model="gpt-4o-mini", temperature=0.2)

# ---------------------------------------------------------------------------
# Node: draft_review
# ---------------------------------------------------------------------------
async def draft_review(state: CodeReviewState) -> CodeReviewState:
    """Generate an initial code review.

    The review contains 3–6 bullet points describing what is good and what can be
    improved.  The output is plain text.
    """
    prompt = (
        "You are a senior Python developer.\n"
        "Given the following function, write a concise code review (3–6 points).\n"
        "Focus on style, correctness, and potential improvements.\n"
        "Return only the review text.\n\n"
        f"Function:\n{state["code"]}\n"
    )
    review = await llm.ainvoke(prompt)
    state["draft_review"] = review.content.strip()
    return state

# ---------------------------------------------------------------------------
# Node: reflect
# ---------------------------------------------------------------------------
class ReflectOutput(BaseModel):
    pep8: int = Field(..., ge=0, le=10)
    type_hints: int = Field(..., ge=0, le=10)
    edge_cases: int = Field(..., ge=0, le=10)
    naming: int = Field(..., ge=0, le=10)
    verdict: str = Field(..., regex="^(ok|needs_revision)$")
    weakest_criterion: str = Field(..., regex="^(pep8|type_hints|edge_cases|naming)$")

async def reflect(state: CodeReviewState) -> CodeReviewState:
    """Critic node that scores the draft review on four criteria.

    The LLM returns a JSON object that matches ``ReflectOutput``.  The function
    parses the JSON and updates the state.
    """
    prompt = (
        "You are a code quality critic.\n"
        "Given the following code review, score it on the following criteria (0–10):\n"
        "- PEP8 compliance\n"
        "- Type hints usage\n"
        "- Edge cases handling\n"
        "- Naming conventions\n"
        "Return a JSON object with keys: pep8, type_hints, edge_cases, naming, verdict, weakest_criterion.\n"
        "Verdict should be "ok" if all scores are >=7, otherwise "needs_revision".\n"
        "Weakest criterion is the one with the lowest score.\n\n"
        f"Review:\n{state["draft_review"]}\n"
    )
    result = await llm.ainvoke(prompt)
    try:
        data = ReflectOutput.model_validate_json(result.content)
    except Exception as e:
        # Fallback: if parsing fails, treat as needs_revision
        data = ReflectOutput(
            pep8=0,
            type_hints=0,
            edge_cases=0,
            naming=0,
            verdict="needs_revision",
            weakest_criterion="pep8",
        )
    state["criteria_scores"] = data.model_dump(exclude="verdict,weakest_criterion")
    state["weakest_criterion"] = data.weakest_criterion
    state["verdict"] = data.verdict
    return state

# ---------------------------------------------------------------------------
# Node: rewrite
# ---------------------------------------------------------------------------
async def rewrite(state: CodeReviewState) -> CodeReviewState:
    """Rewrite the weakest part of the review.

    The LLM is instructed to rewrite only the section that addresses the
    weakest criterion.  The new review replaces the old one.
    """
    prompt = (
        "You are a senior Python developer.\n"
        "Rewrite the part of the following code review that addresses the weakest criterion.\n"
        "Keep the rest of the review unchanged.\n"
        "Return only the updated review text.\n\n"
        f"Weakest criterion: {state["weakest_criterion"]}\n"
        f"Current review:\n{state["draft_review"]}\n"
    )
    new_review = await llm.ainvoke(prompt)
    state["draft_review"] = new_review.content.strip()
    state["round"] += 1
    return state

# ---------------------------------------------------------------------------
# Graph construction
# ---------------------------------------------------------------------------
def build_graph() -> StateGraph[CodeReviewState]:
    graph = StateGraph(CodeReviewState)
    graph.add_node("draft_review", draft_review)
    graph.add_node("reflect", reflect)
    graph.add_node("rewrite", rewrite)

    # Entry point
    graph.set_entry_point("draft_review")

    # Transitions
    graph.add_conditional_edges(
        "draft_review",
        lambda state: "reflect",
    )
    graph.add_conditional_edges(
        "reflect",
        lambda state: "rewrite" if state["verdict"] == "needs_revision" and state["round"] < state["max_rounds"] else "END",
    )
    graph.add_edge("rewrite", "reflect")

    return graph

# ---------------------------------------------------------------------------
# Demo helper
# ---------------------------------------------------------------------------
async def run_demo(func) -> None:
    """Run the graph on a single function and print the results."""
    import inspect
    import asyncio

    code = inspect.getsource(func)
    state: CodeReviewState = {
        "code": code,
        "draft_review": None,
        "criteria_scores": None,
        "weakest_criterion": None,
        "verdict": None,
        "round": 0,
        "max_rounds": 2,
    }
    graph = build_graph()
    final_state = await graph.astate(state)
    print("\n=== Initial Review ===")
    print(state["draft_review"])
    print("\n=== Scores ===")
    print(final_state["criteria_scores"])
    print("Verdict:", final_state["verdict"])
    if final_state["verdict"] == "needs_revision":
        print("\n=== Revised Review ===")
        print(final_state["draft_review"])

# ---------------------------------------------------------------------------
# Example function for demo
# ---------------------------------------------------------------------------

def sort_numbers(arr):
    """Return a sorted copy of the input list."""
    return sorted(arr)

# ---------------------------------------------------------------------------
# Main entry point
# ---------------------------------------------------------------------------
if __name__ == "__main__":
    import asyncio
    asyncio.run(run_demo(sort_numbers))