"""
# main.py
# LangGraph code review agent with deepagents integration
# Author: OpenAI ChatGPT
# Requirements: deepagents, langchain>=1.2.10, langchain-openai>=0.3.0, langgraph>=0.2.0
# Run: python main.py
"""

import os
import json
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

# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------
# Load OpenRouter API key from environment
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
if not OPENAI_API_KEY:
    raise RuntimeError("OPENAI_API_KEY environment variable not set")

# LLM instance (OpenRouter)
llm = ChatOpenAI(
    model="openai/gpt-oss-20b:free",
    base_url="https://openrouter.ai/api/v1",
    api_key=OPENAI_API_KEY,
    temperature=0.0,
)

# ---------------------------------------------------------------------------
# State definition
# ---------------------------------------------------------------------------
class CodeReviewState(TypedDict):
    code: str
    draft_review: str
    criteria_scores: Dict[str, int]  # {"pep8": int, "type_hints": int, "edge_cases": int, "naming": int}
    weakest_criterion: str
    verdict: str  # "ok" | "needs_revision"
    round: int
    max_rounds: int

# ---------------------------------------------------------------------------
# Pydantic model for reflect output
# ---------------------------------------------------------------------------
class ReviewScores(BaseModel):
    pep8: int = Field(..., description="Score for PEP8 compliance (0-10)")
    type_hints: int = Field(..., description="Score for type hints usage (0-10)")
    edge_cases: int = Field(..., description="Score for handling edge cases (0-10)")
    naming: int = Field(..., description="Score for naming conventions (0-10)")
    verdict: str = Field(..., description="'ok' if all scores >=7 else 'needs_revision'")

# ---------------------------------------------------------------------------
# LangGraph nodes
# ---------------------------------------------------------------------------
async def draft_review(state: CodeReviewState) -> CodeReviewState:
    code = state["code"]
    prompt = f"""
Write a concise code review (3-6 bullet points) for the following Python function:

{code}

Review:
"""
    messages = [HumanMessage(content=prompt)]
    response = await llm.ainvoke(messages)
    review = response.content.strip()
    state["draft_review"] = review
    return state

async def reflect(state: CodeReviewState) -> CodeReviewState:
    review = state["draft_review"]
    code = state["code"]
    prompt = f"""
You are a code review critic. Evaluate the following review against the code. Assign scores 0-10 for each criterion: PEP8, type hints, edge cases, naming. Also provide verdict: "ok" if all scores >=7, else "needs_revision". Return JSON with keys: pep8, type_hints, edge_cases, naming, verdict.

Example:
{{"pep8":8,"type_hints":9,"edge_cases":6,"naming":7,"verdict":"needs_revision"}}

Code:
{code}

Review:
{review}

JSON:
"""
    messages = [HumanMessage(content=prompt)]
    response = await llm.ainvoke(messages)
    try:
        data = json.loads(response.content)
    except Exception as e:
        data = {}
    scores = {
        "pep8": int(data.get("pep8", 0)),
        "type_hints": int(data.get("type_hints", 0)),
        "edge_cases": int(data.get("edge_cases", 0)),
        "naming": int(data.get("naming", 0)),
    }
    weakest = min(scores, key=scores.get)
    state["criteria_scores"] = scores
    state["weakest_criterion"] = weakest
    state["verdict"] = data.get("verdict", "needs_revision")
    return state

async def rewrite(state: CodeReviewState) -> CodeReviewState:
    review = state["draft_review"]
    weakest = state["weakest_criterion"]
    code = state["code"]
    prompt = f"""
Rewrite the part of the review that addresses the {weakest} criterion to improve it. Keep other parts unchanged. Provide only the updated review.

Original review:
{review}

Updated review:
"""
    messages = [HumanMessage(content=prompt)]
    response = await llm.ainvoke(messages)
    updated_review = response.content.strip()
    state["draft_review"] = updated_review
    state["round"] += 1
    return state

# ---------------------------------------------------------------------------
# Build the graph
# ---------------------------------------------------------------------------
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")
# After draft_review always go to reflect
graph.add_edge("draft_review", "reflect")
# Conditional after reflect
def reflect_cond(state: CodeReviewState):
    if state["verdict"] == "needs_revision" and state["round"] < state["max_rounds"]:
        return "rewrite"
    return "END"

graph.add_conditional_edges("reflect", reflect_cond)
# After rewrite go back to reflect
graph.add_edge("rewrite", "reflect")

compiled_graph = graph.compile()

# ---------------------------------------------------------------------------
# DeepAgents integration
# ---------------------------------------------------------------------------
# Backend for deepagents (in-memory shell, no real files used)
backend = CompositeBackend(
    default=LocalShellBackend(root_dir="./workspace", virtual_mode=True, inherit_env=True),
    routes={},
)

@tool
def run_code_review(code: str) -> str:
    """Run a multi‑round code review on the provided Python function."""
    # Initial state
    state: CodeReviewState = {
        "code": code,
        "draft_review": "",
        "criteria_scores": {},
        "weakest_criterion": "",
        "verdict": "",
        "round": 0,
        "max_rounds": 2,
    }

    async def _run():
        final_state = await compiled_graph.invoke(state)
        # Logging
        print("\n--- Draft Review ---")
        print(final_state["draft_review"])
        print("\n--- Scores ---")
        print(final_state["criteria_scores"])
        if final_state["verdict"] == "needs_revision":
            print("\n--- Rewrite performed ---")
        print("\n--- Final Review ---")
        print(final_state["draft_review"])
        return final_state["draft_review"]

    return asyncio.run(_run())

agent = create_deep_agent(
    model=llm,
    tools=[run_code_review],
    backend=backend,
    system_prompt="You are a helpful code review assistant. Use the provided tool to review code.",
)

# ---------------------------------------------------------------------------
# Demo CLI
# ---------------------------------------------------------------------------
if __name__ == "__main__":
    sample_code = """
# Example function to sort numbers

def sort_numbers(arr):
    return sorted(arr)
"""
    # Directly invoke the tool for demonstration
    print("Running demo on sample function...")
    final_review = run_code_review(sample_code)
    print("\n=== Final Review Output ===")
    print(final_review)