import os
import asyncio
from typing import TypedDict, Annotated

from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage, AIMessage
from langchain_core.output_parsers import PydanticOutputParser
from langchain_core.pydantic_v1 import BaseModel, Field
from langgraph.graph import StateGraph, START, END
from langgraph.graph.message import add_messages

# ---------- LLM ----------
llm = ChatOpenAI(
    model="openai/gpt-oss-20b:free",
    base_url="https://openrouter.ai/api/v1",
    api_key=os.getenv("OPENAI_API_KEY"),
    temperature=0.0,
)

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

# ---------- Structured output for reflect ----------
class ReflectionOutput(BaseModel):
    pep8: int = Field(description="Score 0-10 for PEP8 compliance")
    type_hints: int = Field(description="Score 0-10 for type hints usage")
    edge_cases: int = Field(description="Score 0-10 for handling edge cases")
    naming: int = Field(description="Score 0-10 for naming conventions")
    weakest_criterion: str = Field(description="Criterion with lowest score")
    verdict: str = Field(description="'ok' or 'needs_revision'")

parser = PydanticOutputParser(pydantic_object=ReflectionOutput)

# ---------- Nodes ----------

def draft_review_node(state: CodeReviewState) -> CodeReviewState:
    code = state["code"]
    prompt = (
        "You are a senior Python developer.\n"
        "Given the following function, write a concise code review (3-6 bullet points).\n"
        "Focus on style, correctness, edge cases, and naming.\n"
        f"Function:\n{code}\n\nReview:"  # LLM will output review
    )
    response = llm.invoke([HumanMessage(content=prompt)])
    state["draft_review"] = response.content
    return state


def reflect_node(state: CodeReviewState) -> CodeReviewState:
    review = state["draft_review"]
    code = state["code"]
    prompt = (
        "You are an automated code review critic.\n"
        "Given the code and its review, assign a score 0-10 for each of the following criteria:\n"
        "- pep8: PEP8 compliance\n"
        "- type_hints: use of type hints\n"
        "- edge_cases: handling of edge cases\n"
        "- naming: clarity of names\n"
        "Return the scores, the weakest criterion, and a verdict ('ok' if all scores >=7, else 'needs_revision').\n"
        f"Code:\n{code}\n\nReview:\n{review}\n\nOutput in JSON with fields: pep8, type_hints, edge_cases, naming, weakest_criterion, verdict."
    )
    response = llm.invoke([HumanMessage(content=prompt)])
    try:
        out = parser.parse(response.content)
    except Exception as e:
        # Fallback: simple parsing if JSON is not strict
        import json
        out = json.loads(response.content)
    state["criteria_scores"] = {
        "pep8": out.pep8,
        "type_hints": out.type_hints,
        "edge_cases": out.edge_cases,
        "naming": out.naming,
    }
    state["weakest_criterion"] = out.weakest_criterion
    state["verdict"] = out.verdict
    return state


def rewrite_node(state: CodeReviewState) -> CodeReviewState:
    weakest = state["weakest_criterion"]
    review = state["draft_review"]
    code = state["code"]
    prompt = (
        "You are a senior Python developer tasked with improving a code review.\n"
        f"The current review is:\n{review}\n\nThe weakest criterion is '{weakest}'.\n"
        "Rewrite only the part of the review that addresses this criterion, making it stronger and more specific.\n"
        "Keep the rest of the review unchanged.\n"
        "Output only the updated review."
    )
    response = llm.invoke([HumanMessage(content=prompt)])
    state["draft_review"] = response.content
    state["round"] = state.get("round", 0) + 1
    return state

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

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")
builder.add_edge("END", END)

graph = builder.compile()

# ---------- Demo ----------
async def main():
    demo_code = """
    def sort_numbers(arr):
        return sorted(arr)
    """
    initial_state: CodeReviewState = {
        "code": demo_code.strip(),
        "draft_review": "",  # will be filled
        "criteria_scores": {},
        "weakest_criterion": "",
        "verdict": "",
        "round": 0,
        "max_rounds": 2,
    }
    result = await graph.ainvoke(initial_state)
    print("\n--- Final Review ---")
    print(result["draft_review"])
    print("\n--- Scores ---")
    for k, v in result["criteria_scores"].items():
        print(f"{k}: {v}")
    print(f"Verdict: {result['verdict']}")

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