"""
# main.py
# LangGraph code review agent with reflection and rewrite loop
# Requires: langgraph, langchain-openai, deepagents, python-dotenv
# Run with: python main.py
"""

from __future__ import annotations

import os
from typing import TypedDict, Dict, Any

from dotenv import load_dotenv

# Deepagents is required by the assignment, but we do not use it directly.
# Importing it ensures that the dependency is satisfied.
import deepagents  # noqa: F401

from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate, SystemMessagePromptTemplate, HumanMessagePromptTemplate
from langchain_core.output_parsers import PydanticOutputParser
from langchain_core.messages import HumanMessage
from langgraph.graph import StateGraph, END
from langgraph.checkpoint.memory import MemorySaver
from pydantic import BaseModel

# Load OpenAI key from .env if present
load_dotenv()

# -----------------------------
# State definition
# -----------------------------
class CodeReviewState(TypedDict):
    code: str
    draft_review: str
    criteria_scores: Dict[str, int]  # e.g., {'pep8': 8, ...}
    weakest_criterion: str
    verdict: str  # "ok" | "needs_revision"
    round: int
    max_rounds: int

# -----------------------------
# LLM configuration
# -----------------------------
# Replace with your preferred model or use Ollama via langchain-ollama if desired
llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)

# -----------------------------
# Prompt templates
# -----------------------------
# draft_review prompt
draft_prompt = ChatPromptTemplate.from_messages([
    SystemMessagePromptTemplate.from_template(
        "You are a senior Python developer. Provide a concise code review for the following function. "
        "List 3-6 bullet points, each describing a potential improvement or praise."
    ),
    HumanMessagePromptTemplate.from_template("Here is the function code:\n\n{code}")
])

# reflect prompt with structured output
class ReviewScores(BaseModel):
    pep8: int
    type_hints: int
    edge_cases: int
    naming: int
    weakest_criterion: str
    verdict: str

output_parser = PydanticOutputParser(pydantic_object=ReviewScores)

reflect_prompt = ChatPromptTemplate.from_messages([
    SystemMessagePromptTemplate.from_template(
        "You are a code quality critic. Evaluate the draft review for the following function. "
        "Score each of the four criteria (pep8, type_hints, edge_cases, naming) on a scale 0-10. "
        "Identify the weakest criterion and provide a verdict: 'ok' if all scores are >=7, otherwise 'needs_revision'."
    ),
    HumanMessagePromptTemplate.from_template("Function code:\n\n{code}\n\nDraft review:\n\n{draft_review}")
])

# rewrite prompt
rewrite_prompt = ChatPromptTemplate.from_messages([
    SystemMessagePromptTemplate.from_template(
        "You are a senior Python developer. Rewrite the draft review to improve the weakest criterion: {weakest_criterion}. "
        "Keep 3-6 bullet points."
    ),
    HumanMessagePromptTemplate.from_template("Original draft review:\n\n{draft_review}")
])

# -----------------------------
# Node implementations
# -----------------------------
async def draft_review(state: CodeReviewState) -> Dict[str, Any]:
    """Generate initial draft review."""
    response = await llm.ainvoke(draft_prompt.format_messages(code=state['code']))
    draft = response.content.strip()
    # Ensure 3-6 bullet points by counting lines starting with '-'
    bullets = [line for line in draft.splitlines() if line.lstrip().startswith('-')]
    if len(bullets) < 3:
        # If too few, add a generic positive point
        draft += "\n- The code is readable and concise."
    elif len(bullets) > 6:
        # Trim to 6
        draft = "\n".join(bullets[:6])
    return {"draft_review": draft}

async def reflect(state: CodeReviewState) -> Dict[str, Any]:
    """Critic evaluates the draft review."""
    response = await llm.ainvoke(
        reflect_prompt.format_messages(
            code=state['code'],
            draft_review=state['draft_review']
        )
    )
    parsed = output_parser.parse(response.content)
    return {
        "criteria_scores": {
            "pep8": parsed.pep8,
            "type_hints": parsed.type_hints,
            "edge_cases": parsed.edge_cases,
            "naming": parsed.naming
        },
        "weakest_criterion": parsed.weakest_criterion,
        "verdict": parsed.verdict
    }

async def rewrite(state: CodeReviewState) -> Dict[str, Any]:
    """Rewrite the draft review focusing on the weakest criterion."""
    response = await llm.ainvoke(
        rewrite_prompt.format_messages(
            weakest_criterion=state['weakest_criterion'],
            draft_review=state['draft_review']
        )
    )
    new_draft = response.content.strip()
    # Ensure 3-6 bullet points
    bullets = [line for line in new_draft.splitlines() if line.lstrip().startswith('-')]
    if len(bullets) < 3:
        new_draft += "\n- The code is readable and concise."
    elif len(bullets) > 6:
        new_draft = "\n".join(bullets[:6])
    return {"draft_review": new_draft, "round": state['round'] + 1}

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

    # Entry point
    graph.set_entry_point("draft_review")

    # Transitions
    graph.add_edge("draft_review", "reflect")
    graph.add_conditional_edges(
        "reflect",
        lambda x: "needs_revision" if x["verdict"] == "needs_revision" else "ok",
        {
            "needs_revision": "rewrite",
            "ok": END
        }
    )
    graph.add_edge("rewrite", "reflect")

    # Final node to stop if max rounds reached
    def check_round(state: CodeReviewState):
        if state["round"] >= state["max_rounds"]:
            return "END"
        return "rewrite"

    graph.add_conditional_edges(
        "rewrite",
        check_round,
        {
            "END": END,
            "rewrite": "rewrite"
        }
    )

    return graph.compile(checkpointer=MemorySaver())

# -----------------------------
# Demo execution
# -----------------------------
if __name__ == "__main__":
    # Sample function to review
    sample_code = """
    def sort_numbers(arr):
        return sorted(arr)
    """

    # Initial state
    state: CodeReviewState = {
        "code": sample_code.strip(),
        "draft_review": "",
        "criteria_scores": {},
        "weakest_criterion": "",
        "verdict": "",
        "round": 0,
        "max_rounds": 2
    }

    graph = build_graph()

    # Run the graph
    async def run():
        async for event in graph.astream(state):
            # Print only the updated parts for clarity
            if "draft_review" in event:
                print("\n--- Draft Review (Round %d) ---" % (event.get("round", 0)))
                print(event["draft_review"])
            if "criteria_scores" in event:
                print("\n--- Scores ---")
                for k, v in event["criteria_scores"].items():
                    print(f"{k}: {v}")
                print(f"Weakest criterion: {event['weakest_criterion']}")
                print(f"Verdict: {event['verdict']}")

    import asyncio
    asyncio.run(run())