import os
import asyncio
from typing import TypedDict, Annotated, Dict
from langgraph.graph import StateGraph, START, END
from langgraph.graph.message import add_messages
from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage, AIMessage
from langchain_core.output_parsers import PydanticOutputParser
from pydantic import BaseModel, Field
from deepagents import create_deep_agent
from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend
from deepagents.tools import tool

# ---------- 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
    round: int
    max_rounds: int

# ---------- Structured output for reflect ----------
class ReflectOutput(BaseModel):
    pep8: int = Field(..., ge=0, le=10)
    type_hints: int = Field(..., ge=0, le=10)
    edge_cases: int = Field(..., ge=0, le=10)
    naming: int = Field(..., ge=0, le=10)
    weakest_criterion: str = Field(...)
    verdict: str = Field(..., regex="^(ok|needs_revision)$")

reflect_parser = PydanticOutputParser(pydantic_object=ReflectOutput)

# ---------- Nodes ----------
@tool
def draft_review(state: CodeReviewState) -> CodeReviewState:
    """Generate an initial code review."""
    prompt = (
        "You are a senior Python reviewer.\n"
        "Given the following function, write a concise code review (3–6 points).\n"
        "Focus on style, correctness, and potential improvements.\n"
        "Return only the review text.\n\n"
        f"Function:\n{state['code']}"
    )
    review = llm.invoke([HumanMessage(content=prompt)]).content
    state['draft_review'] = review
    return state

@tool
def reflect(state: CodeReviewState) -> CodeReviewState:
    """Critique the draft review and score four criteria."""
    prompt = (
        "You are an automated code review critic.\n"
        "Given the original code and the draft review, assign a score 0–10 for each of the following criteria:\n"
        "- pep8: adherence to PEP8 style guide\n"
        "- type_hints: use of type hints\n"
        "- edge_cases: handling of edge cases\n"
        "- naming: clarity of identifiers\n"
        "Also identify the weakest criterion and decide if the review is "ok" or "needs_revision".\n"
        "Return a JSON object with keys: pep8, type_hints, edge_cases, naming, weakest_criterion, verdict.\n"
        "Do not include any other text.\n\n"
        f"Code:\n{state['code']}\n\n"
        f"Draft Review:\n{state['draft_review']}"
    )
    raw = llm.invoke([HumanMessage(content=prompt)]).content
    try:
        parsed = reflect_parser.parse(raw)
    except Exception as e:
        # Fallback: simple extraction
        parsed = ReflectOutput(pep8=5, type_hints=5, edge_cases=5, naming=5, weakest_criterion="pep8", verdict="needs_revision")
    state['criteria_scores'] = {
        "pep8": parsed.pep8,
        "type_hints": parsed.type_hints,
        "edge_cases": parsed.edge_cases,
        "naming": parsed.naming,
    }
    state['weakest_criterion'] = parsed.weakest_criterion
    state['verdict'] = parsed.verdict
    return state

@tool
def rewrite(state: CodeReviewState) -> CodeReviewState:
    """Rewrite the section of the draft review that addresses the weakest criterion."""
    prompt = (
        "You are a senior Python reviewer.\n"
        "The draft review below has been critiqued. The weakest criterion is {criterion}.\n"
        "Rewrite only the part of the review that addresses this criterion, improving it.\n"
        "Keep the rest of the review unchanged.\n"
        "Return the full updated review.\n\n"
        f"Weakest criterion: {state['weakest_criterion']}\n\n"
        f"Draft Review:\n{state['draft_review']}"
    ).format(criterion=state['weakest_criterion'])
    updated = llm.invoke([HumanMessage(content=prompt)]).content
    state['draft_review'] = updated
    state['round'] += 1
    return state

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

graph = builder.compile()

# ---------- DeepAgent wrapper ----------
backend = CompositeBackend([
    LocalShellBackend(workspace_dir="./workspace"),
    FilesystemBackend(),
])

@tool
def run_review(code: str) -> str:
    """Run the LangGraph code review pipeline on the provided code."""
    initial_state: CodeReviewState = {
        "code": code,
        "draft_review": "",
        "criteria_scores": {},
        "weakest_criterion": "",
        "verdict": "",
        "round": 0,
        "max_rounds": 2,
    }
    final_state = graph.invoke(initial_state)
    return (
        f"Initial Draft Review:\n{final_state['draft_review']}\n\n"
        f"Scores: {final_state['criteria_scores']}\n"
        f"Verdict: {final_state['verdict']}\n"
        f"Rounds: {final_state['round']}\n"
    )

agent = create_deep_agent(
    model=llm,
    tools=[run_review],
    backend=backend,
    system_prompt="You are a code review assistant.",
)

async def main():
    code_example = """
    def sort_numbers(arr):
        return sorted(arr)
    """
    result = await agent.ainvoke(
        {"messages": [HumanMessage(content=f"Please review this code:\n{code_example}")]},
        {"configurable": {"thread_id": "session-1"}},
    )
    print(result["messages"][-1].content)

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