import asyncio
import os
from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage
from langchain.tools import tool
from langgraph.graph import StateGraph, START, END
from typing import TypedDict, Annotated
from langgraph.graph.message import add_messages
from pydantic import BaseModel, Field
from langchain_core.output_parsers import PydanticOutputParser

# LLM setup
llm = ChatOpenAI(
    model="openai/gpt-4o-mini",
    base_url="https://openrouter.ai/api/v1",
    api_key=os.getenv("OPENAI_API_KEY"),
    temperature=0.0,
)

# Pydantic models for reflection
class CriteriaScores(BaseModel):
    pep8: int
    type_hints: int
    edge_cases: int
    naming: int

class ReviewResult(BaseModel):
    scores: CriteriaScores
    weakest: str
    verdict: str

# Tool to run shell commands
@tool
def run_command(command: str) -> str:
    """Execute a shell command and return its output."""
    try:
        result = os.popen(command).read()
        return result.strip() or "(no output)"
    except Exception as e:
        return f"Error: {e}"

# Draft review node
async def draft_review(state: dict):
    code = state["code"]
    prompt = f"Write a concise code review for the following Python function. Provide 3-6 bullet points highlighting strengths and areas for improvement.\n\n{code}"
    response = llm.invoke([HumanMessage(content=prompt)])
    state["draft_review"] = response.content
    return state

# Reflect node
async def reflect(state: dict):
    review = state["draft_review"]
    prompt = f"Score the following code review on 4 criteria: PEP8, type hints, edge cases, naming. Return JSON with keys pep8, type_hints, edge_cases, naming (0-10). Also provide the weakest criterion and verdict ('ok' if all >=7 else 'needs_revision').\n\n{review}"
    response = llm.invoke([HumanMessage(content=prompt)])
    try:
        data = ReviewResult.parse_raw(response.content)
    except Exception:
        # fallback simple parse
        data = ReviewResult(scores=CriteriaScores(pep8=5,type_hints=5,edge_cases=5,naming=5),weakest="pep8",verdict="needs_revision")
    state["criteria_scores"] = data.scores.dict()
    state["weakest_criterion"] = data.weakest
    state["verdict"] = data.verdict
    return state

# Rewrite node
async def rewrite(state: dict):
    weakest = state["weakest_criterion"]
    review = state["draft_review"]
    prompt = f"Improve the code review focusing on the {weakest} aspect. Keep the rest unchanged.\n\n{review}"
    response = llm.invoke([HumanMessage(content=prompt)])
    state["draft_review"] = response.content
    state["round"] += 1
    return state

# Graph definition
class CodeReviewState(TypedDict):
    code: str
    draft_review: str
    criteria_scores: dict
    weakest_criterion: str
    verdict: str
    round: int
    max_rounds: int

workflow = StateGraph(CodeReviewState)
workflow.add_node("draft", draft_review)
workflow.add_node("reflect", reflect)
workflow.add_node("rewrite", rewrite)
workflow.add_conditional_edges(START, lambda _: "draft")
workflow.add_conditional_edges("draft", lambda _: "reflect")
workflow.add_conditional_edges("reflect", lambda s: "rewrite" if s["verdict"]=="needs_revision" and s["round"]<s["max_rounds"] else "END")
workflow.add_conditional_edges("rewrite", lambda _: "reflect")
graph = workflow.compile()

# Demo function
async def run_demo():
    code = """
def sort_numbers(arr):
    return sorted(arr)
"""
    state: CodeReviewState = {
        "code": code,
        "draft_review": "",
        "criteria_scores": {},
        "weakest_criterion": "",
        "verdict": "",
        "round": 0,
        "max_rounds": 2,
    }
    final = await graph.ainvoke(state)
    print("Final review:\n", final["draft_review"])
    print("Scores:", final["criteria_scores"])

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