# main.py
import os
from typing import TypedDict, Dict
from langgraph.graph import StateGraph, END
from langgraph.prebuilt import create_chat_agent
from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage, AIMessage

# Define 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
llm = ChatOpenAI(temperature=0)

# Draft review node
async def draft_review(state: CodeReviewState) -> CodeReviewState:
    prompt = f"""
You are a senior Python developer. Review the following code and provide a concise code review (3-6 bullet points) highlighting what is good and what can be improved.

Code:
{state['code']}

Review:
"""
    response = await llm.ainvoke([HumanMessage(content=prompt)])
    state['draft_review'] = response.content
    return state

# Reflect node
async def reflect(state: CodeReviewState) -> CodeReviewState:
    prompt = f"""
You are an AI critic evaluating a code review. Assign a score 0-10 for each of the following criteria based on the draft review:
- pep8
- type_hints
- edge_cases
- naming

Provide a JSON object with keys "pep8", "type_hints", "edge_cases", "naming" and integer values.
Also determine the weakest criterion (the one with lowest score) and a verdict: "ok" if all scores >=7, otherwise "needs_revision".

Draft review:
{state['draft_review']}

Output JSON:
"""
    response = await llm.ainvoke([HumanMessage(content=prompt)])
    import json
    scores = json.loads(response.content)
    state['criteria_scores'] = scores
    weakest = min(scores, key=scores.get)
    state['weakest_criterion'] = weakest
    state['verdict'] = "ok" if all(v >= 7 for v in scores.values()) else "needs_revision"
    return state

# Rewrite node
async def rewrite(state: CodeReviewState) -> CodeReviewState:
    crit = state['weakest_criterion']
    prompt = f"""
You are a senior Python developer. Rewrite the section of the code review that addresses the {crit} criterion, improving it. Keep the rest of the review unchanged.

Original review:
{state['draft_review']}

Rewrite only the part related to {crit}:
"""
    response = await llm.ainvoke([HumanMessage(content=prompt)])
    # Replace the part in draft_review that mentions crit
    # For simplicity, just append the new part
    state['draft_review'] = state['draft_review'] + "\n" + response.content
    state['round'] += 1
    return state

# Build 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: "rewrite" if x["verdict"] == "needs_revision" and x["round"] < x["max_rounds"] else "END",
)
builder.add_edge("rewrite", "reflect")

graph = builder.compile()

# Demo function
async def run_demo():
    code = """
# Example function to sort numbers

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("Final Review:\n", result["draft_review"])
    print("Scores:\n", result["criteria_scores"])
    print("Verdict:\n", result["verdict"])

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