from typing import Annotated
from typing_extensions import TypedDict
from langgraph.graph import StateGraph, START, END
from langgraph.graph.message import add_messages
from dotenv import load_dotenv
from langgraph.prebuilt import ToolNode
from langchain_openai import ChatOpenAI
from langgraph.checkpoint.memory import MemorySaver
from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
from typing import List, Any, Optional, Dict
from pydantic import BaseModel, Field
from sidekick_tools import playwright_tools, other_tools
import uuid
import asyncio
from datetime import datetime

load_dotenv(override=True)


class State(TypedDict):
    messages: Annotated[List[Any], add_messages]
    success_criteria: str
    feedback_on_work: Optional[str]
    success_criteria_met: bool
    user_input_needed: bool
    clarifying_questions: List[str]   # NEW: 3 questions generated by clarifier
    clarifying_answers: List[str]     # NEW: answers provided by the user


class EvaluatorOutput(BaseModel):
    feedback: str = Field(description="Feedback on the assistant's response")
    success_criteria_met: bool = Field(description="Whether the success criteria have been met")
    user_input_needed: bool = Field(
        description="True if more input is needed from the user, or clarifications, or the assistant is stuck"
    )


class ClarifyingQuestions(BaseModel):
    questions: List[str] = Field(
        description="Exactly 3 clarifying questions to ask the user before starting the task"
    )


class Sidekick:
    def __init__(self):
        self.worker_llm_with_tools = None
        self.evaluator_llm_with_output = None
        self.clarifier_llm_with_output = None   # NEW
        self.tools = None
        self.llm_with_tools = None
        self.graph = None
        self.sidekick_id = str(uuid.uuid4())
        self.memory = MemorySaver()
        self.browser = None
        self.playwright = None

    async def setup(self):
        self.tools, self.browser, self.playwright = await playwright_tools()
        self.tools += await other_tools()

        worker_llm = ChatOpenAI(model="gpt-4o-mini")
        self.worker_llm_with_tools = worker_llm.bind_tools(self.tools)

        evaluator_llm = ChatOpenAI(model="gpt-4o-mini")
        self.evaluator_llm_with_output = evaluator_llm.with_structured_output(EvaluatorOutput)

        # NEW: clarifier LLM with structured output
        clarifier_llm = ChatOpenAI(model="gpt-4o-mini")
        self.clarifier_llm_with_output = clarifier_llm.with_structured_output(ClarifyingQuestions)

        await self.build_graph()

    # ── NEW NODE ──────────────────────────────────────────────

    def clarifier(self, state: State) -> Dict[str, Any]:
        """
        Runs once at the start of every new task.
        Generates 3 targeted clarifying questions based on the task
        and success criteria, then pauses for user answers.
        Skipped entirely if questions have already been asked.
        """
        # Guard — only ask once per task
        if state.get("clarifying_questions"):
            return {}

        last_message = state["messages"][-1].content

        prompt = f"""You are about to complete a task for a user using browser tools, 
search, file management, and Python execution.

Before starting, generate exactly 3 clarifying questions that would most 
help you complete the task successfully and avoid wrong assumptions.

Task: {last_message}
Success criteria: {state['success_criteria']}

Focus your questions on:
- Scope or constraints (e.g. which website, which format, date range)
- Output preferences (e.g. detailed or summary, file or text)
- Any ambiguity that would cause you to guess wrong

Be specific and concise. Do not ask generic questions."""

        result = self.clarifier_llm_with_output.invoke(
            [HumanMessage(content=prompt)]
        )

        questions_text = "\n".join(
            f"{i+1}. {q}" for i, q in enumerate(result.questions)
        )

        return {
            "clarifying_questions": result.questions,
            "messages": [AIMessage(
                content=f"Before I begin, I have 3 quick questions:\n\n{questions_text}"
            )],
        }

    def clarifier_router(self, state: State) -> str:
        """
        After clarifier runs:
        - Questions just generated and not yet answered → END (pause for user)
        - Questions already answered → proceed to worker
        """
        questions = state.get("clarifying_questions", [])
        answers   = state.get("clarifying_answers",   [])

        if questions and len(answers) < len(questions):
            return END       # pause — Gradio collects answers
        return "worker"      # answers in state — start working

    # ── EXISTING NODES (unchanged except worker system prompt) ─

    def worker(self, state: State) -> Dict[str, Any]:

        # Build clarification context from Q&A pairs
        clarification_context = ""
        if state.get("clarifying_questions") and state.get("clarifying_answers"):
            pairs = zip(state["clarifying_questions"], state["clarifying_answers"])
            clarification_context = "\n\nClarifications provided by the user:\n" + \
                "\n".join(f"Q: {q}\nA: {a}" for q, a in pairs)

        system_message = f"""You are a helpful assistant that can use tools to complete tasks.
You keep working on a task until either you have a question or clarification for the user, or the success criteria is met.
You have many tools to help you, including tools to browse the internet, navigating and retrieving web pages.
You have a tool to run python code, but note that you would need to include a print() statement if you wanted to receive output.
The current date and time is {datetime.now().strftime("%Y-%m-%d %H:%M:%S")}

This is the success criteria:
{state["success_criteria"]}
{clarification_context}

You should reply either with a question for the user about this assignment, or with your final response.
If you have a question for the user, you need to reply by clearly stating your question. An example might be:

Question: please clarify whether you want a summary or a detailed answer

If you've finished, reply with the final answer, and don't ask a question; simply reply with the answer.
"""

        if state.get("feedback_on_work"):
            system_message += f"""
Previously you thought you completed the assignment, but your reply was rejected because the success criteria was not met.
Here is the feedback on why this was rejected:
{state["feedback_on_work"]}
With this feedback, please continue the assignment, ensuring that you meet the success criteria or have a question for the user."""

        # Update or insert system message
        found_system_message = False
        messages = state["messages"]
        for message in messages:
            if isinstance(message, SystemMessage):
                message.content = system_message
                found_system_message = True

        if not found_system_message:
            messages = [SystemMessage(content=system_message)] + messages

        response = self.worker_llm_with_tools.invoke(messages)
        return {"messages": [response]}

    def worker_router(self, state: State) -> str:
        last_message = state["messages"][-1]
        if hasattr(last_message, "tool_calls") and last_message.tool_calls:
            return "tools"
        else:
            return "evaluator"

    def format_conversation(self, messages: List[Any]) -> str:
        conversation = "Conversation history:\n\n"
        for message in messages:
            if isinstance(message, HumanMessage):
                conversation += f"User: {message.content}\n"
            elif isinstance(message, AIMessage):
                text = message.content or "[Tools use]"
                conversation += f"Assistant: {text}\n"
        return conversation

    def evaluator(self, state: State) -> State:
        last_response = state["messages"][-1].content

        system_message = """You are an evaluator that determines if a task has been completed successfully by an Assistant.
Assess the Assistant's last response based on the given criteria. Respond with your feedback, and with your decision on whether the success criteria has been met,
and whether more input is needed from the user."""

        user_message = f"""You are evaluating a conversation between the User and Assistant. You decide what action to take based on the last response from the Assistant.

The entire conversation with the assistant, with the user's original request and all replies, is:
{self.format_conversation(state["messages"])}

The success criteria for this assignment is:
{state["success_criteria"]}

And the final response from the Assistant that you are evaluating is:
{last_response}

Respond with your feedback, and decide if the success criteria is met by this response.
Also, decide if more user input is required, either because the assistant has a question, needs clarification, or seems to be stuck and unable to answer without help.

The Assistant has access to a tool to write files. If the Assistant says they have written a file, then you can assume they have done so.
Overall you should give the Assistant the benefit of the doubt if they say they've done something. But you should reject if you feel that more work should go into this.
"""
        if state["feedback_on_work"]:
            user_message += f"Also, note that in a prior attempt from the Assistant, you provided this feedback: {state['feedback_on_work']}\n"
            user_message += "If you're seeing the Assistant repeating the same mistakes, then consider responding that user input is required."

        evaluator_messages = [
            SystemMessage(content=system_message),
            HumanMessage(content=user_message),
        ]

        eval_result = self.evaluator_llm_with_output.invoke(evaluator_messages)
        return {
            "messages": [{
                "role": "assistant",
                "content": f"Evaluator Feedback on this answer: {eval_result.feedback}",
            }],
            "feedback_on_work":      eval_result.feedback,
            "success_criteria_met":  eval_result.success_criteria_met,
            "user_input_needed":     eval_result.user_input_needed,
        }

    def route_based_on_evaluation(self, state: State) -> str:
        if state["success_criteria_met"] or state["user_input_needed"]:
            return "END"
        else:
            return "worker"

    # ── GRAPH BUILD (updated with clarifier) ──────────────────

    async def build_graph(self):
        graph_builder = StateGraph(State)

        # Nodes
        graph_builder.add_node("clarifier", self.clarifier)   # NEW
        graph_builder.add_node("worker",    self.worker)
        graph_builder.add_node("tools",     ToolNode(tools=self.tools))
        graph_builder.add_node("evaluator", self.evaluator)

        # Edges
        graph_builder.add_edge(START, "clarifier")             # NEW entry point
        graph_builder.add_conditional_edges(                   # NEW
            "clarifier", self.clarifier_router,
            {"worker": "worker", END: END}
        )
        graph_builder.add_conditional_edges(
            "worker", self.worker_router,
            {"tools": "tools", "evaluator": "evaluator"}
        )
        graph_builder.add_edge("tools", "worker")
        graph_builder.add_conditional_edges(
            "evaluator", self.route_based_on_evaluation,
            {"worker": "worker", "END": END}
        )

        self.graph = graph_builder.compile(checkpointer=self.memory)

    # ── RUN (updated to carry clarifying answers) ─────────────

    async def run_superstep(self, message, success_criteria, history,
                             clarifying_answers: List[str] = None):
        config = {"configurable": {"thread_id": self.sidekick_id}}

        state = {
            "messages":            message,
            "success_criteria":    success_criteria or "The answer should be clear and accurate",
            "feedback_on_work":    None,
            "success_criteria_met": False,
            "user_input_needed":   False,
            "clarifying_questions": [],
            "clarifying_answers":  clarifying_answers or [],  # NEW
        }

        result = await self.graph.ainvoke(state, config=config)

        # Check if we paused for clarifying questions
        questions = result.get("clarifying_questions", [])
        answers   = result.get("clarifying_answers",   [])

        if questions and len(answers) < len(questions):
            # Graph paused — return questions for Gradio to display
            last_msg = result["messages"][-1].content
            return history + [
                {"role": "user",      "content": message},
                {"role": "assistant", "content": last_msg},
            ], questions

        # Graph completed — return normal worker + evaluator reply
        user     = {"role": "user",      "content": message}
        reply    = {"role": "assistant", "content": result["messages"][-2].content}
        feedback = {"role": "assistant", "content": result["messages"][-1].content}
        return history + [user, reply, feedback], []

    def cleanup(self):
        if self.browser:
            try:
                loop = asyncio.get_running_loop()
                loop.create_task(self.browser.close())
                if self.playwright:
                    loop.create_task(self.playwright.stop())
            except RuntimeError:
                asyncio.run(self.browser.close())
                if self.playwright:
                    asyncio.run(self.playwright.stop())