import os
import asyncio
from typing import TypedDict, Annotated, Dict

from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage
from langchain.tools import tool

from langgraph.graph import StateGraph, START, END
from langgraph.graph.message import add_messages
from pydantic import BaseModel, Field
from langchain_core.output_parsers import PydanticOutputParser

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

# ---------- Pydantic for reflect output ----------
class ReflectOutput(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 edge case handling")
    naming: int = Field(..., description="Score 0-10 for naming conventions")
    weakest_criterion: str = Field(..., description="Name of the weakest criterion")
    verdict: str = Field(..., description="'ok' or 'needs_revision'")

reflect_parser = PydanticOutputParser(pydantic_object=ReflectOutput)

# ---------- Nodes ----------
async def draft_review(state: CodeReviewState) -> CodeReviewState:
    prompt = f"""Please write a concise code review (3-6 bullet points) for the following Python function. Focus on style, type hints, edge cases, and naming.

```python
{state['code']}
```

Return only the review text."""
    review = await llm.ainvoke([HumanMessage(content=prompt)])
    state['draft_review'] = review.content.strip()
    return state

async def reflect(state: CodeReviewState) -> CodeReviewState:
    prompt = f"""You are a code review critic. Evaluate the following review text and assign scores 0-10 for each of the four criteria: pep8, type_hints, edge_cases, naming. Also identify the weakest criterion and decide if the review is "ok" or "needs_revision".

Review text:
{state['draft_review']}

Provide the output in the following JSON-like format:
{"pep8": int, "type_hints": int, "edge_cases": int, "naming": int, "weakest_criterion": str, "verdict": str}"""
    raw = await llm.ainvoke([HumanMessage(content=prompt)])
    parsed = reflect_parser.parse(raw.content)
    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

async def rewrite(state: CodeReviewState) -> CodeReviewState:
    # Simple rewrite: add a sentence addressing the weakest criterion
    additional = f"Additionally, the review should pay more attention to {state['weakest_criterion']}."
    state['draft_review'] = state['draft_review'] + "\n" + additional
    state['round'] += 1
    return state

# ---------- Graph ----------
def build_graph() -> StateGraph[CodeReviewState]:
    graph = StateGraph(CodeGraphState)
    graph.add_node("draft_review", draft_review)
    graph.add_node("reflect", reflect)
    graph.add_node("rewrite", rewrite)

    graph.set_entry_point("draft_review")
    graph.add_edge("draft_review", "reflect")
    graph.add_conditional_edges(
        "reflect",
        lambda x: "END" if x['verdict'] == "ok" or x['round'] >= x['max_rounds'] else "rewrite",
    )
    graph.add_edge("rewrite", "reflect")

    return graph.compile()

# ---------- Tool ----------
@tool
def code_review_tool(code: str) -> str:
    """Perform a structured code review with possible rewrites."""
    graph = build_graph()
    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"Final Review:\n{final_state['draft_review']}\n\nScores: {final_state['criteria_scores']}"

# ---------- DeepAgent ----------
async def main():
    sample_code = """
    def sort_numbers(arr):
        return sorted(arr)
    """
    # Use the code_review_tool directly without deepagents
    final_review = code_review_tool(sample_code)
    print(final_review)

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