"""
LangGraph Code Review Agent
Author: Auto-generated
"""
import os
from typing import TypedDict, Dict
from langgraph.graph import StateGraph
# from langgraph.prebuilt import create_agent_executor  # unused
from langchain_openai import ChatOpenAI
from dotenv import load_dotenv

load_dotenv()

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

# ---------- Nodes ----------
async def draft_review(state: CodeReviewState) -> CodeReviewState:
    llm = ChatOpenAI(model="gpt-4o-mini", temperature=0.2)
    prompt = (
        "You are a senior Python developer.\n"
        f"Given the following code:\n{state['code']}\n"
        "Write a concise code review (3–6 points). Do not include any other text."  # noqa: E501
    )
    response = await llm.invoke(prompt)
    state["draft_review"] = response.content.strip()
    return state

async def reflect(state: CodeReviewState) -> CodeReviewState:
    llm = ChatOpenAI(model="gpt-4o-mini", temperature=0.1)
    prompt = f'''You are an automated code review critic.
Draft review:
{state['draft_review']}
Score the draft on the following criteria (0–10):
- PEP8 compliance
- Type hints usage
- Edge case handling
- Naming conventions
Return a JSON object with keys: pep8, type_hints, edge_cases, naming, weakest_criterion, verdict. The verdict should be "ok" if all scores are >=7, otherwise "needs_revision".'''
    response = await llm.invoke(prompt)
    import json
    try:
        data = json.loads(response.content.strip())
    except Exception as e:
        raise ValueError(f"Failed to parse critic output: {e}\n{response.content}")
    state["criteria_scores"] = {
        "pep8": int(data.get("pep8", 0)),
        "type_hints": int(data.get("type_hints", 0)),
        "edge_cases": int(data.get("edge_cases", 0)),
        "naming": int(data.get("naming", 0)),
    }
    state["weakest_criterion"] = data.get("weakest_criterion", "")
    state["verdict"] = data.get("verdict", "needs_revision")
    return state

async def rewrite(state: CodeReviewState) -> CodeReviewState:
    llm = ChatOpenAI(model="gpt-4o-mini", temperature=0.2)
    prompt = (
        f"The draft review was insufficient in the {state['weakest_criterion']} area.\n"
        "Rewrite only that part of the review to improve it, keeping other points unchanged.\n"
        f"Current draft:\n{state['draft_review']}\n"
        "Provide the updated draft."  # noqa: E501
    )
    response = await llm.invoke(prompt)
    state["draft_review"] = response.content.strip()
    state["round"] += 1
    return state

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

builder.set_entry_point("draft")
builder.add_edge("draft", "reflect")
builder.add_conditional_edges(
    "reflect",
    lambda state: (
        "end"
        if state["verdict"] == "ok"
        else ("rewrite" if state["round"] < state.get("max_rounds", 2) else "end")
    ),
)
builder.add_edge("rewrite", "reflect")

graph = builder.compile()

# ---------- Demo ----------
if __name__ == "__main__":
    sample_code = (
        "def sort_numbers(arr):\n"
        "    return sorted(arr)"
    )
    initial_state: CodeReviewState = {
        "code": sample_code,
        "draft_review": "",
        "criteria_scores": {},
        "weakest_criterion": "",
        "verdict": "needs_revision",
        "round": 0,
        "max_rounds": 2,
    }
    result = graph.invoke(initial_state)
    print("\n--- Final Review ---")
    print(result["draft_review"])
    print("\nScores:", result["criteria_scores"])
    print("Verdict:", result["verdict"])