import os
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
from langchain_core.output_parsers import PydanticOutputParser

# ---------- 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):
    pep8: int = Field(description="Score for PEP8 compliance (0-10)")
    type_hints: int = Field(description="Score for type hints (0-10)")
    edge_cases: int = Field(description="Score for edge case handling (0-10)")
    naming: int = Field(description="Score for naming conventions (0-10)")
    verdict: str = Field(description='Verdict: "ok" or "needs_revision"')

critic_parser = PydanticOutputParser(pydantic_object=CriticOutput)

# ---------- LLM and backend ----------
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 = CompositeBackend(
    [
        LocalShellBackend(workspace_dir="./workspace"),
        FilesystemBackend(),
    ]
)

agent = create_deep_agent(
    model=llm,
    tools=[],
    backend=backend,
    system_prompt="You are a helpful agent.",
)

# ---------- Node functions ----------
async def draft_review(state: CodeReviewState) -> CodeReviewState:
    prompt = f"Please provide a concise code review (3-6 points) for the following Python function:\n\n{state['code']}"
    result = await agent.ainvoke([HumanMessage(content=prompt)])
    review = result["messages"][-1].content.strip()
    state["draft_review"] = review
    print("\n--- Draft Review ---")
    print(review)
    return state

async def reflect(state: CodeReviewState) -> CodeReviewState:
    prompt = (
        f"Evaluate the following draft review:\n\n{state['draft_review']}\n\n"
        "Score each of the following criteria on a scale of 0-10:\n"
        "- pep8\n- type_hints\n- edge_cases\n- naming\n\n"
        "Return the scores and a verdict ('ok' or 'needs_revision') in the following JSON format:\n"
        "{\n  \"pep8\": int,\n  \"type_hints\": int,\n  \"edge_cases\": int,\n  \"naming\": int,\n  \"verdict\": \"ok\" | \"needs_revision\"\n}"
    )
    result = await agent.ainvoke([HumanMessage(content=prompt)])
    raw_output = result["messages"][-1].content.strip()
    try:
        parsed = critic_parser.parse(raw_output)
    except Exception as e:
        # Fallback: simple parsing if LLM output is not perfectly formatted
        parsed = CriticOutput(
            pep8=0,
            type_hints=0,
            edge_cases=0,
            naming=0,
            verdict="needs_revision",
        )
    scores = {
        "pep8": parsed.pep8,
        "type_hints": parsed.type_hints,
        "edge_cases": parsed.edge_cases,
        "naming": parsed.naming,
    }
    weakest = min(scores, key=scores.get)
    state["criteria_scores"] = scores
    state["weakest_criterion"] = weakest
    state["verdict"] = parsed.verdict
    print("\n--- Reflection ---")
    print(f"Scores: {scores}")
    print(f"Weakest criterion: {weakest}")
    print(f"Verdict: {parsed.verdict}")
    return state

async def rewrite(state: CodeReviewState) -> CodeReviewState:
    prompt = (
        f"Rewrite the section of the draft review that addresses the weakest criterion "
        f"('{state['weakest_criterion']}') to improve it. Keep all other parts unchanged.\n\n"
        f"Original draft review:\n\n{state['draft_review']}"
    )
    result = await agent.ainvoke([HumanMessage(content=prompt)])
    new_review = result["messages"][-1].content.strip()
    state["draft_review"] = new_review
    state["round"] += 1
    print("\n--- Rewritten Review ---")
    print(new_review)
    return state

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

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

    def reflect_cond(state: CodeReviewState):
        if state["verdict"] == "ok":
            return END
        if state["round"] < state["max_rounds"]:
            return "rewrite"
        return END

    graph.add_conditional_edges("reflect", reflect_cond, {"rewrite": "rewrite", END: END})
    graph.add_edge("rewrite", "reflect")

    return graph

# ---------- Demo ----------
async def main():
    # Sample function to review
    code_str = """def sort_numbers(arr):
    return sorted(arr)"""

    initial_state: CodeReviewState = {
        "code": code_str,
        "draft_review": "",
        "criteria_scores": {},
        "weakest_criterion": "",
        "verdict": "",
        "round": 0,
        "max_rounds": 2,
    }

    graph = build_graph()
    app = graph.compile()
    final_state = await app.ainvoke(initial_state)

    print("\n=== Final State ===")
    print(f"Round: {final_state['round']}")
    print(f"Verdict: {final_state['verdict']}")
    print(f"Draft Review:\n{final_state['draft_review']}")
    print(f"Scores: {final_state['criteria_scores']}")

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