"""
LangGraph Code Review Agent
===========================

This repository contains a minimal LangGraph implementation that
performs a code review on a Python function.  The graph consists of
three nodes:

* ``draft_review`` – generates an initial review.
* ``reflect`` – a critic that scores the review on four criteria
  (PEP8, type hints, edge cases, naming) and decides whether a
  rewrite is required.
* ``rewrite`` – rewrites the weakest part of the review.

The graph runs for a maximum of ``max_rounds`` (default 2).  The
demo can be executed with ``python -m main``.
"""

from __future__ import annotations

import os
from typing import TypedDict, Dict

from langgraph.graph import StateGraph, END
from langgraph.prebuilt import create_react_agent
from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage

# ---------------------------------------------------------------------------
# State definition
# ---------------------------------------------------------------------------

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

# ---------------------------------------------------------------------------
# LLM configuration
# ---------------------------------------------------------------------------

# The user must set the OPENAI_API_KEY environment variable.
llm = ChatOpenAI(temperature=0.0, model="gpt-4o-mini")

# ---------------------------------------------------------------------------
# Node implementations
# ---------------------------------------------------------------------------

async def draft_review(state: CodeReviewState) -> CodeReviewState:
    """Generate an initial review of the provided code.

    The review is a short list of 3–6 bullet points.
    """
    code = state["code"]
    prompt = (
        "You are a senior Python developer.\n"
        "Review the following function and provide a concise list of 3–6 points\n"
        "highlighting what is good and what could be improved.\n"
        "Do not mention the criteria – just give the review.\n"
        f"Function:\n{code}\n"
        "Review:"  # LLM will continue after this
    )
    review = await llm.ainvoke([HumanMessage(content=prompt)])
    state["draft_review"] = review.content.strip()
    return state

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

    The output is a JSON object with keys:
    * pep8, type_hints, edge_cases, naming – integers 0‑10
    * weakest_criterion – one of the four keys
    * verdict – "ok" or "needs_revision"
    """
    review = state["draft_review"]
    code = state["code"]
    prompt = (
        "You are a code quality critic.\n"
        "Given the following code and its draft review, score the review on\n"
        "four criteria: PEP8, type hints, edge cases, naming.\n"
        "Return a JSON object with keys: pep8, type_hints, edge_cases, naming,\n"
        "weakest_criterion, verdict.\n"
        f"Code:\n{code}\n"
        f"Draft review:\n{review}\n"
        "Answer in JSON only."
    )
    response = await llm.ainvoke([HumanMessage(content=prompt)])
    import json
    try:
        scores = json.loads(response.content)
    except Exception as e:
        # Fallback: if parsing fails, treat as needs_revision
        scores = {
            "pep8": 0,
            "type_hints": 0,
            "edge_cases": 0,
            "naming": 0,
            "weakest_criterion": "pep8",
            "verdict": "needs_revision",
        }
    state["criteria_scores"] = {
        "pep8": int(scores.get("pep8", 0)),
        "type_hints": int(scores.get("type_hints", 0)),
        "edge_cases": int(scores.get("edge_cases", 0)),
        "naming": int(scores.get("naming", 0)),
    }
    state["weakest_criterion"] = scores.get("weakest_criterion", "pep8")
    state["verdict"] = scores.get("verdict", "needs_revision")
    return state

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

    The node receives the current state and the weakest criterion.
    It generates a new review that specifically addresses that criterion.
    """
    code = state["code"]
    weakest = state["weakest_criterion"]
    prompt = (
        "You are a senior Python developer.\n"
        "Rewrite the draft review to improve the part related to the following criterion: "
        f"{weakest}.\n"
        "Keep the rest of the review unchanged.\n"
        "Output only the updated review.\n"
        f"Current draft review:\n{state['draft_review']}\n"
    )
    new_review = await llm.ainvoke([HumanMessage(content=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_edge("draft_review", "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 execution
# ---------------------------------------------------------------------------

if __name__ == "__main__":
    import argparse
    import textwrap

    parser = argparse.ArgumentParser(description="Run the code review graph on a demo function.")
    parser.add_argument("--max-rounds", type=int, default=2, help="Maximum number of rewrite rounds")
    args = parser.parse_args()

    # Demo function – can be replaced by any user code
    demo_code = textwrap.dedent(
        """
        def sort_numbers(arr):
            return sorted(arr)
        """
    ).strip()

    initial_state: CodeReviewState = {
        "code": demo_code,
        "draft_review": "",
        "criteria_scores": {},
        "weakest_criterion": "",
        "verdict": "",
        "round": 0,
        "max_rounds": args.max_rounds,
    }

    graph = build_graph()
    final_state = graph.invoke(initial_state)

    print("\n=== Final Review ===")
    print(final_state["draft_review"])
    print("\n=== Scores ===")
    for k, v in final_state["criteria_scores"].items():
        print(f"{k}: {v}")
    print(f"Verdict: {final_state['verdict']}")
    print(f"Rounds: {final_state['round']}")