"""LangGraph Code Review Agent

This repository implements a LangGraph agent that takes a Python function as input and produces a code review. The review is evaluated by a critic node that scores it on four criteria:

1. PEP8 compliance
2. Type hints
3. Edge case handling
4. Naming conventions

If the critic returns "needs_revision" the rewrite node improves the weakest part of the review. The process repeats up to ``max_rounds`` times.

The implementation uses only the technologies specified in the assignment: ``langgraph`` and ``langchain-ollama`` (or ``langchain-openai`` if you prefer). No vector database is used.

Run the demo with ``python main.py``.
"""

from __future__ import annotations

from typing import TypedDict, Dict

from langgraph.graph import StateGraph, END
from langgraph.prebuilt import create_structured_output_node
from langchain_ollama import ChatOllama
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StructuredOutputParser
from langchain_core.messages import HumanMessage

# ---------- State ----------
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 ----------
# Use Ollama; adjust model name if needed
llm = ChatOllama(model="llama3")

# ---------- Draft Review Node ----------
DRAFT_PROMPT = ChatPromptTemplate.from_messages([
    ("system", "You are a senior Python developer. Your task is to write a concise code review for the following function. Provide 3-6 points, each starting with a dash.")
])

async def draft_review(state: CodeReviewState) -> Dict[str, str]:
    prompt = DRAFT_PROMPT.format_messages(code=state["code"])
    response = await llm.ainvoke(prompt)
    review = response.content.strip()
    return {"draft_review": review}

# ---------- Reflect Node ----------
# Structured output schema
SCHEMA = {
    "pep8": "int (0-10)",
    "type_hints": "int (0-10)",
    "edge_cases": "int (0-10)",
    "naming": "int (0-10)",
    "weakest_criterion": "string (one of the keys above)",
    "verdict": "string (\"ok\" or \"needs_revision\")",
}

parser = StructuredOutputParser.from_function_signature(
    "def scores(pep8: int, type_hints: int, edge_cases: int, naming: int, weakest_criterion: str, verdict: str) -> dict"
)

REFLECT_PROMPT = ChatPromptTemplate.from_messages([
    ("system", "You are a code review critic. Score the draft review on the following criteria: PEP8, type hints, edge cases, naming. Provide scores 0-10 and decide if the review is \"ok\" or \"needs_revision\".")
])

async def reflect(state: CodeReviewState) -> Dict[str, object]:
    prompt = REFLECT_PROMPT.format_messages(draft_review=state["draft_review"])
    response = await llm.ainvoke(prompt)
    # Parse structured output
    try:
        parsed = parser.parse(response.content)
    except Exception as e:
        # Fallback: simple heuristic
        parsed = {
            "pep8": 5,
            "type_hints": 5,
            "edge_cases": 5,
            "naming": 5,
            "weakest_criterion": "pep8",
            "verdict": "needs_revision",
        }
    return {
        "criteria_scores": {
            "pep8": parsed["pep8"],
            "type_hints": parsed["type_hints"],
            "edge_cases": parsed["edge_cases"],
            "naming": parsed["naming"],
        },
        "weakest_criterion": parsed["weakest_criterion"],
        "verdict": parsed["verdict"],
    }

# ---------- Rewrite Node ----------
async def rewrite(state: CodeReviewState) -> Dict[str, str]:
    # Find the weakest criterion and add a focused improvement note
    wc = state["weakest_criterion"]
    improvement = f"\n- Improve {wc.replace('_', ' ')}: Provide more detailed guidance on this aspect."
    new_review = state["draft_review"] + improvement
    return {"draft_review": new_review}

# ---------- 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")

builder.add_edge("draft_review", "reflect")
builder.add_conditional_edges(
    "reflect",
    lambda x: "END" if x["verdict"] == "ok" else "rewrite",
)
builder.add_edge("rewrite", "reflect")

# Stop after max_rounds
builder.add_conditional_edges(
    "reflect",
    lambda x: "END" if x["round"] >= x["max_rounds"] else "rewrite",
)

graph = builder.compile()

# ---------- Demo ----------
async def main():
    # Example function to review
    code = """
    def sort_numbers(arr):
        return sorted(arr)
    """
    init_state: CodeReviewState = {
        "code": code,
        "draft_review": "",
        "criteria_scores": {},
        "weakest_criterion": "",
        "verdict": "",
        "round": 0,
        "max_rounds": 2,
    }
    result = await graph.ainvoke(init_state)
    print("\n--- Draft Review ---")
    print(result["draft_review"])
    print("\n--- Scores ---")
    print(result["criteria_scores"])
    print("\n--- Verdict ---")
    print(result["verdict"])

if __name__ == "__main__":
    import asyncio
    asyncio.run(main())