import os
import asyncio
from typing import TypedDict, 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 pydantic import BaseModel, Field
from langchain_core.output_parsers import PydanticOutputParser

# LLM setup
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 definition
class CodeReviewState(TypedDict):
    code: str
    draft_review: str
    criteria_scores: Dict[str, int]
    weakest_criterion: str
    verdict: str
    round: int
    max_rounds: int

# Reflect output model
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 coverage")
    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'")

reflect_parser = PydanticOutputParser(pydantic_object=ReflectOutput)

# Dummy tool for agent
@tool
def echo_tool(query: str) -> str:
    return query

# Agent creation
from deepagents import create_deep_agent
from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend
backend = CompositeBackend(
    default=LocalShellBackend(root_dir="./workspace", virtual_mode=True, inherit_env=True),
    routes={},
)
agent = create_deep_agent(
    model=llm,
    tools=[echo_tool],
    backend=backend,
    system_prompt="You are a code review assistant.",
)

# Node functions
async def draft_review(state: CodeReviewState) -> CodeReviewState:
    prompt = f"""Write a concise code review (3-6 points) for the following Python function. Focus on style, correctness, and potential improvements.

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

Return only the review text."""
    response = await agent.ainvoke({"messages": [HumanMessage(content=prompt)]}, {"configurable": {"thread_id": "draft-review"}})
    review_text = response["messages"][-1].content
    state["draft_review"] = review_text
    return state

async def reflect(state: CodeReviewState) -> CodeReviewState:
    prompt = f"""You are a senior reviewer. Evaluate the following review text against four criteria: PEP8, type hints, edge cases, naming. Assign each a score 0-10. Identify the weakest criterion and give a verdict: 'ok' if all scores >=7, else 'needs_revision'. Return a JSON with keys: pep8, type_hints, edge_cases, naming, weakest_criterion, verdict.

Review:
{state['draft_review']}"""
    response = await agent.ainvoke({"messages": [HumanMessage(content=prompt)]}, {"configurable": {"thread_id": "reflect"}})
    json_text = response["messages"][-1].content
    try:
        parsed = reflect_parser.parse(json_text)
    except Exception:
        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

async def rewrite(state: CodeReviewState) -> CodeReviewState:
    crit = state["weakest_criterion"]
    prompt = f"""Improve the review section that addresses the weakest criterion '{crit}'. Provide a more detailed point for that criterion. Keep the rest of the review unchanged.

Current review:
{state['draft_review']}"""
    response = await agent.ainvoke({"messages": [HumanMessage(content=prompt)]}, {"configurable": {"thread_id": "rewrite"}})
    new_review = response["messages"][-1].content
    state["draft_review"] = new_review
    state["round"] += 1
    return state

# Graph definition
from langgraph.graph import StateGraph

graph = StateGraph(CodeReviewState)

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" else "rewrite" if x["round"] < x["max_rounds"] else "END",
)

graph.add_edge("rewrite", "reflect")
app = graph.compile()

# Demo function
async def demo():
    sample_code = """def sort_numbers(arr):
    return sorted(arr)"""
    init_state: CodeReviewState = {
        "code": sample_code,
        "draft_review": "",
        "criteria_scores": {},
        "weakest_criterion": "",
        "verdict": "",
        "round": 0,
        "max_rounds": 2,
    }
    result = await app.ainvoke(init_state)
    print("--- Draft Review ---")
    print(result["draft_review"])
    print("\n--- Scores ---")
    print(result["criteria_scores"])
    print("\n--- Verdict ---")
    print(result["verdict"])
    if result["verdict"] == "needs_revision":
        print("\n--- Final Review After Rewrite ---")
        print(result["draft_review"])
        print("\n--- Final Scores ---")
        print(result["criteria_scores"])

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