import os
import asyncio
from typing import TypedDict, Dict

from dotenv import load_dotenv
from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage
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
from deepagents import create_deep_agent
from deepagents.backends import FilesystemBackend

# Load API key from .env
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 - simple filesystem
backend = FilesystemBackend()

# Create a deepagents agent that will be used inside the graph nodes
agent = create_deep_agent(
    model=llm,
    tools=[],
    backend=backend,
    system_prompt="You are a code review assistant.",
)

# ---------- 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 critic ----------
class CriticOutput(BaseModel):
    scores: Dict[str, int] = Field(
        description="Scores for each criterion: pep8, type_hints, edge_cases, naming. Values 0-10."
    )
    verdict: str = Field(
        description='Verdict: "ok" if all scores >= 7, otherwise "needs_revision".'
    )

critic_parser = PydanticOutputParser(pydantic_object=CriticOutput)

# ---------- Graph nodes ----------
async def draft_review(state: CodeReviewState) -> CodeReviewState:
    prompt = (
        f"Write a concise code review (3-6 bullet points) for the following Python function:\n\n"
        f"{state['code']}\n\n"
        "Focus on style, correctness, and potential improvements."
    )
    response = await agent.ainvoke(
        {"messages": [HumanMessage(content=prompt)]},
        {"configurable": {"thread_id": "draft_review"}},
    )
    review_text = response["messages"][-1].content.strip()
    state["draft_review"] = review_text
    print("\n--- Draft Review ---")
    print(review_text)
    return state

async def reflect(state: CodeReviewState) -> CodeReviewState:
    prompt = (
        f"Evaluate the following code review and assign scores (0-10) for each criterion:\n\n"
        f"Review:\n{state['draft_review']}\n\n"
        "Criteria:\n"
        "1. pep8: adherence to PEP8 style guide.\n"
        "2. type_hints: presence and correctness of type hints.\n"
        "3. edge_cases: handling of edge cases and robustness.\n"
        "4. naming: clarity and consistency of names.\n\n"
        "Return a JSON object with keys 'scores' (dict) and 'verdict' ('ok' or 'needs_revision')."
    )
    response = await agent.ainvoke(
        {"messages": [HumanMessage(content=prompt)]},
        {"configurable": {"thread_id": "reflect"}},
    )
    raw_output = response["messages"][-1].content.strip()
    try:
        parsed = critic_parser.parse(raw_output)
    except Exception as e:
        # Fallback: simple parsing if JSON is malformed
        import json
        parsed = CriticOutput(**json.loads(raw_output))
    state["criteria_scores"] = parsed.scores
    # Determine weakest criterion
    weakest = min(parsed.scores.items(), key=lambda kv: kv[1])[0]
    state["weakest_criterion"] = weakest
    state["verdict"] = parsed.verdict
    print("\n--- Critic Scores ---")
    for crit, score in parsed.scores.items():
        print(f"{crit}: {score}")
    print(f"Weakest criterion: {weakest}")
    print(f"Verdict: {parsed.verdict}")
    return state

async def rewrite(state: CodeReviewState) -> CodeReviewState:
    state["round"] += 1
    prompt = (
        f"Rewrite the part of the review that addresses the weakest criterion "
        f"('{state['weakest_criterion']}') to improve it. Keep the rest of the review unchanged.\n\n"
        f"Original Review:\n{state['draft_review']}\n\n"
        "Provide only the updated review."
    )
    response = await agent.ainvoke(
        {"messages": [HumanMessage(content=prompt)]},
        {"configurable": {"thread_id": "rewrite"}},
    )
    new_review = response["messages"][-1].content.strip()
    state["draft_review"] = new_review
    print("\n--- Rewritten Review (Round {}) ---".format(state["round"]))
    print(new_review)
    return state

# ---------- Graph construction ----------
builder = StateGraph(CodeReviewState)

builder.add_node("draft_review", draft_review)
builder.add_node("reflect", reflect)
builder.add_node("rewrite", rewrite)

builder.add_edge(START, "draft_review")
builder.add_edge("draft_review", "reflect")

# Conditional edges after reflect
def reflect_conditional(state: CodeReviewState):
    if state["verdict"] == "ok":
        return END
    if state["round"] < state["max_rounds"]:
        return "rewrite"
    return END

builder.add_conditional_edges("reflect", reflect_conditional)

builder.add_edge("rewrite", "reflect")

graph = builder.compile()

# ---------- Demo ----------
async def main():
    # Sample function to review
    sample_code = """
def sort_numbers(arr):
    return sorted(arr)
"""
    initial_state: CodeReviewState = {
        "code": sample_code.strip(),
        "draft_review": "",
        "criteria_scores": {},
        "weakest_criterion": "",
        "verdict": "",
        "round": 0,
        "max_rounds": 2,
    }

    final_state = await graph.ainvoke(initial_state)
    print("\n=== Final State ===")
    print(f"Verdict: {final_state['verdict']}")
    print(f"Rounds performed: {final_state['round']}")
    print("\nFinal Review:")
    print(final_state["draft_review"])

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