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 deepagents import create_deep_agent
from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend
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

# Load environment variables
from dotenv import load_dotenv
load_dotenv()

# LLM configuration - OpenRouter
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,
)

# Backend for deepagents
backend = CompositeBackend(
    [
        LocalShellBackend(workspace_dir="./workspace"),
        FilesystemBackend(),
    ]
)

# State definition
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 node
class ReviewScores(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")
    verdict: str = Field(description='\"ok\" or \"needs_revision\"')

review_parser = PydanticOutputParser(pydantic_object=ReviewScores)

# Draft review node
def draft_review(state: CodeReviewState) -> CodeReviewState:
    prompt = f"""Write a concise code review for the following Python function. Provide 3-6 bullet points. Do not include any code blocks. The function is:
```python
{state["code"]}
```"""
    response = llm.invoke(prompt)
    state["draft_review"] = response.content.strip()
    print("\n--- Draft Review ---")
    print(state["draft_review"])
    return state

# Reflect node
def reflect(state: CodeReviewState) -> CodeReviewState:
    prompt = f"""You are a code review critic. Evaluate the following draft review for the given code. Provide scores 0-10 for each criterion and a verdict. Return JSON matching the schema:
{review_parser.get_format_instructions()}
Draft review:
{state["draft_review"]}

Code:
```python
{state["code"]}
```"""
    response = llm.invoke(prompt)
    parsed = review_parser.parse(response.content)
    state["criteria_scores"] = {
        "pep8": parsed.pep8,
        "type_hints": parsed.type_hints,
        "edge_cases": parsed.edge_cases,
        "naming": parsed.naming,
    }
    state["verdict"] = parsed.verdict
    state["weakest_criterion"] = min(state["criteria_scores"], key=state["criteria_scores"].get)
    print("\n--- Reflection ---")
    print(f"Scores: {state['criteria_scores']}")
    print(f"Weakest criterion: {state['weakest_criterion']}")
    print(f"Verdict: {state['verdict']}")
    return state

# Rewrite node
def rewrite(state: CodeReviewState) -> CodeReviewState:
    prompt = f"""Rewrite the section of the draft review that addresses the weakest criterion ({state['weakest_criterion']}) to improve it. Keep other parts unchanged. Return only the updated draft review.
Original draft review:
{state['draft_review']}

Code:
```python
{state['code']}
```"""
    response = llm.invoke(prompt)
    state["draft_review"] = response.content.strip()
    state["round"] += 1
    print("\n--- Rewrite ---")
    print(state["draft_review"])
    return state

# Graph definition
def build_graph() -> 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")

    def condition(state: CodeReviewState):
        if state["verdict"] == "ok":
            return "END"
        if state["round"] >= state["max_rounds"]:
            return "END"
        return "rewrite"

    graph.add_conditional_edges("reflect", condition, ["rewrite", "END"])
    graph.add_edge("rewrite", "reflect")

    return graph

# Tool that runs the graph
@tool
def run_review(code: str) -> str:
    """Run a code review cycle on the provided Python function."""
    initial_state: CodeReviewState = {
        "code": code,
        "draft_review": "",
        "criteria_scores": {},
        "weakest_criterion": "",
        "verdict": "",
        "round": 0,
        "max_rounds": 2,
    }
    graph = build_graph()
    final_state = graph.invoke(initial_state)
    output = f"Final review after {final_state['round']} round(s):\n{final_state['draft_review']}\n\nScores: {final_state['criteria_scores']}\nVerdict: {final_state['verdict']}"
    return output

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

# Demo function
demo_code = """
def sort_numbers(arr):
    return sorted(arr)
"""

async def main():
    result = await agent.ainvoke(
        {"messages": [HumanMessage(content=demo_code)]},
        {"configurable": {"thread_id": "session-1"}},
    )
    print("\n=== Final Output ===")
    print(result["messages"][-1].content)

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