import os
import json
from typing import TypedDict, Dict
from langgraph.graph import StateGraph, END
from langgraph.prebuilt import create_agent
from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage, AIMessage

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

# ---------------------
# 2. LLM setup
# ---------------------
# Use OpenAI or Ollama based on env variable
if os.getenv("USE_OLLAMA", "false").lower() == "true":
    from langchain_ollama import ChatOllama
    llm = ChatOllama(model="llama3", temperature=0.2)
else:
    llm = ChatOpenAI(temperature=0.2, model_name="gpt-4o-mini")

# ---------------------
# 3. Node definitions
# ---------------------

def draft_review_fn(state: CodeReviewState) -> Dict:
    code = state["code"]
    prompt = f"""
You are a senior Python developer. Provide a concise code review (3-6 bullet points) for the following function. Focus on style, correctness, and potential improvements.

Function:
{code}

Review:
"""
    response = llm.invoke([HumanMessage(content=prompt)])
    review = response.content.strip()
    return {"draft_review": review}


def reflect_fn(state: CodeReviewState) -> Dict:
    review = state["draft_review"]
    prompt = f"""
You are an automated code review critic. Evaluate the following code review on four criteria: PEP8 compliance, type hints usage, edge case handling, and naming conventions. Assign each a score from 0 to 10. Also determine the weakest criterion and a verdict: "ok" if all scores are 7 or higher, otherwise "needs_revision".

Review:
{review}

Respond in JSON with keys: "pep8", "type_hints", "edge_cases", "naming", "weakest_criterion", "verdict".
"""
    response = llm.invoke([HumanMessage(content=prompt)])
    try:
        data = json.loads(response.content)
    except Exception:
        # Fallback: simple parsing
        data = {
            "pep8": 5,
            "type_hints": 5,
            "edge_cases": 5,
            "naming": 5,
            "weakest_criterion": "pep8",
            "verdict": "needs_revision"
        }
    return {
        "criteria_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_criterion": data.get("weakest_criterion", "pep8"),
        "verdict": data.get("verdict", "needs_revision")
    }


def rewrite_fn(state: CodeReviewState) -> Dict:
    weakest = state["weakest_criterion"]
    review = state["draft_review"]
    prompt = f"""
You are a senior Python developer. The following code review has been identified as weak in the "{weakest}" criterion. Rewrite only the part of the review that addresses this criterion, improving it significantly. Keep the rest of the review unchanged.

Original Review:
{review}

Rewritten Review:
"""
    response = llm.invoke([HumanMessage(content=prompt)])
    new_review = response.content.strip()
    return {"draft_review": new_review, "round": state["round"] + 1}

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

builder.add_node("draft_review", draft_review_fn)
builder.add_node("reflect", reflect_fn)
builder.add_node("rewrite", rewrite_fn)

# Entry point
builder.set_entry_point("draft_review")

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

# Final graph
graph = builder.compile()

# ---------------------
# 5. Demo CLI
# ---------------------
if __name__ == "__main__":
    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
    }
    result = graph.invoke(initial_state)
    print("\n--- Final State ---")
    print(json.dumps(result, indent=2))