from typing import TypedDict, Dict
import inspect

from langgraph.graph import StateGraph, END
from langchain_ollama import Ollama
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import JsonOutputParser

# Define the 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 instance (Ollama)
llm = Ollama(model="llama3.1")

# Node: draft_review

def draft_review(state: CodeReviewState) -> Dict[str, str]:
    prompt = ChatPromptTemplate.from_messages([
        ("system", "You are a senior Python developer. Write a concise code review for the given function. Provide 3-6 actionable points."),
        ("user", "Here is the function:\n{code}")
    ])
    chain = prompt | llm
    review = chain.invoke({"code": state["code"]})
    return {"draft_review": review}

# Node: reflect

def reflect(state: CodeReviewState) -> Dict[str, object]:
    prompt = ChatPromptTemplate.from_messages([
        ("system", """You are a code quality critic. Score the following review on four criteria: PEP8, type hints, edge cases, naming. Return a JSON with integer scores 0-10, the weakest criterion, and verdict \"ok\" or \"needs_revision\".\n""") ,
        ("user", "Review:\n{draft_review}")
    ])
    parser = JsonOutputParser()
    chain = prompt | llm | parser
    result = chain.invoke({"draft_review": state["draft_review"]})
    # result is a dict
    return {
        "criteria_scores": {
            "pep8": result["pep8"],
            "type_hints": result["type_hints"],
            "edge_cases": result["edge_cases"],
            "naming": result["naming"],
        },
        "weakest_criterion": result["weakest_criterion"],
        "verdict": result["verdict"],
    }

# Node: rewrite

def rewrite(state: CodeReviewState) -> Dict[str, str]:
    # Increment round
    state["round"] += 1
    prompt = ChatPromptTemplate.from_messages([
        ("system", "You are a senior Python developer. Rewrite the review to improve the section about {weakest_criterion}. Keep other points unchanged."),
        ("user", "Original review:\n{draft_review}")
    ])
    chain = prompt | llm
    new_review = chain.invoke({"weakest_criterion": state["weakest_criterion"], "draft_review": state["draft_review"]})
    return {"draft_review": new_review}

# Build the graph
builder = StateGraph(CodeReviewState)
builder.add_node("draft_review", draft_review)
builder.add_node("reflect", reflect)
builder.add_node("rewrite", rewrite)

builder.add_edge("draft_review", "reflect")
# Conditional edge after reflect
builder.add_conditional_edges(
    "reflect",
    lambda state: "END" if state["verdict"] == "ok" else "rewrite",
)
builder.add_edge("rewrite", "reflect")

builder.set_entry_point("draft_review")
builder.set_finish_point("END")

graph = builder.compile()

# Demo
if __name__ == "__main__":
    def sort_numbers(arr):
        return sorted(arr)

    code = inspect.getsource(sort_numbers)
    initial_state: CodeReviewState = {
        "code": code,
        "draft_review": "",
        "criteria_scores": {},
        "weakest_criterion": "",
        "verdict": "",
        "round": 0,
        "max_rounds": 2,
    }
    result = graph.invoke(initial_state)
    print("\n--- Draft Review ---")
    print(result["draft_review"])
    print("\n--- Scores ---")
    print(result["criteria_scores"])
    print("\n--- Verdict ---")
    print(result["verdict"])
    print("\n--- Round ---")
    print(result["round"])