import os
import textwrap
from typing import Literal, TypedDict

from dotenv import load_dotenv
from langchain_openai import ChatOpenAI
from langgraph.graph import END, START, StateGraph
from pydantic import BaseModel, Field


class CodeReviewState(TypedDict):
    code: str
    draft_review: str
    criteria_scores: dict[str, int]
    weakest_criterion: str
    verdict: str
    round: int
    max_rounds: int


class ReflectionResult(BaseModel):
    criteria_scores: dict[str, int] = Field(
        description="Scores from 0 to 10 for pep8, type_hints, edge_cases, naming."
    )
    weakest_criterion: Literal["pep8", "type_hints", "edge_cases", "naming"]
    verdict: Literal["ok", "needs_revision"]


def build_llm() -> ChatOpenAI:
    return ChatOpenAI(
        model=os.getenv("OPENAI_MODEL", "openai/gpt-oss-20b"),
        base_url=os.getenv("OPENAI_BASE_URL"),
        api_key=os.getenv("OPENAI_API_KEY", "dummy"),
        temperature=0,
    )


def draft_review(state: CodeReviewState) -> dict:
    llm = build_llm()
    prompt = textwrap.dedent(
        """
        Ты опытный Python code reviewer.
        Напиши code review по функции ниже.
        Требования:
        - 3-6 конкретных пунктов;
        - оцени сильные стороны и что улучшить;
        - обязательно затронь PEP8, type hints, edge cases и naming, если это уместно;
        - ответ пиши на русском.

        Код:
        {code}
        """
    ).strip()
    response = llm.invoke(prompt.format(code=state["code"]))
    return {"draft_review": response.content.strip()}


def reflect(state: CodeReviewState) -> dict:
    llm = build_llm().with_structured_output(ReflectionResult)
    prompt = textwrap.dedent(
        """
        Ты критик качества code review.
        Оцени review по 4 критериям:
        - pep8
        - type_hints
        - edge_cases
        - naming

        Правила:
        - для каждого критерия выставь integer score от 0 до 10;
        - weakest_criterion — самый слабый критерий;
        - verdict = "ok", если review уже достаточно хорошее;
        - verdict = "needs_revision", если самое слабое место стоит усилить.

        Код:
        {code}

        Review:
        {review}
        """
    ).strip()
    result = llm.invoke(
        prompt.format(code=state["code"], review=state["draft_review"])
    )
    return {
        "criteria_scores": result.criteria_scores,
        "weakest_criterion": result.weakest_criterion,
        "verdict": result.verdict,
    }


def rewrite(state: CodeReviewState) -> dict:
    llm = build_llm()
    prompt = textwrap.dedent(
        """
        Ты улучшаешь уже написанный code review.
        Нужно целенаправленно усилить самое слабое место: {weakest_criterion}.

        Требования:
        - сохрани общий формат краткого review;
        - сделай акцент именно на критерии {weakest_criterion};
        - добавь более точные и полезные замечания;
        - ответ пиши на русском;
        - итог должен остаться в формате 3-6 пунктов.

        Код:
        {code}

        Текущий review:
        {review}
        """
    ).strip()
    response = llm.invoke(
        prompt.format(
            weakest_criterion=state["weakest_criterion"],
            code=state["code"],
            review=state["draft_review"],
        )
    )
    return {
        "draft_review": response.content.strip(),
        "round": state["round"] + 1,
    }


def next_step(state: CodeReviewState) -> str:
    if state["verdict"] == "ok":
        return "finish"
    if state["round"] >= state["max_rounds"]:
        return "finish"
    return "rewrite"


def build_graph():
    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")
    builder.add_conditional_edges(
        "reflect",
        next_step,
        {
            "rewrite": "rewrite",
            "finish": END,
        },
    )
    builder.add_edge("rewrite", "reflect")
    return builder.compile()


def demo_code() -> str:
    return textwrap.dedent(
        """
        def sort_numbers(arr):
            return sorted(arr)
        """
    ).strip()


def run_demo() -> None:
    graph = build_graph()
    initial_state: CodeReviewState = {
        "code": demo_code(),
        "draft_review": "",
        "criteria_scores": {},
        "weakest_criterion": "",
        "verdict": "",
        "round": 0,
        "max_rounds": 2,
    }

    print("=== Code Review Reflection Demo ===")
    print("Code under review:")
    print(initial_state["code"])

    final_state = graph.run(initial_state)

    print()
    print("=== Final Result ===")
    print(final_state["draft_review"])
    print(final_state["criteria_scores"])
    print(f"Rounds used: {final_state['round']} / {final_state['max_rounds']}")


if __name__ == "__main__":
    load_dotenv()
    run_demo()