# The core pipeline: plan -> execute (with tools) -> judge -> format, with the
# judge and formatter each able to loop back to the planner with feedback if the
# answer isn't good enough (bounded by MAX_RETRIES).
# LlamaIndex Workflow: an event-driven state machine. Each @step method declares,
# via its type hints, which Event type it consumes and which it produces; the
# framework wires them together into a graph and runs a step whenever a matching
# event arrives. Context is the per-run storage (GaiaState) that survives across
# steps and across loop iterations, since Events themselves are thrown away after
# the step that consumed them returns.
from llama_index.core.agent.workflow import FunctionAgent
from llama_index.core.agent.workflow.workflow_events import ToolCallResult
from llama_index.core.prompts import PromptTemplate
from llama_index.core.workflow import Context, Event, StartEvent, StopEvent, Workflow, step

from gaia_agent.config import MAX_RETRIES
from gaia_agent.formatting import format_answer
from gaia_agent.judge import judge_answer
from gaia_agent.llm import get_llm
from gaia_agent.models import AttemptRecord, ConciseAnswer, ExecutionResult, GaiaState, Plan
from gaia_agent.tools.audio import transcribe_audio_tool
from gaia_agent.tools.file_reader import read_attached_file_tool
from gaia_agent.tools.math_tool import calculate_tool
from gaia_agent.tools.multimodal import describe_image_tool, ocr_tool
from gaia_agent.tools.scraper import fetch_url_text_tool
from gaia_agent.tools.web_search import web_search_tool
from gaia_agent.tools.you_answer import you_answer_tool

EXECUTOR_TOOLS = [
    calculate_tool,
    you_answer_tool,
    web_search_tool,
    fetch_url_text_tool,
    describe_image_tool,
    ocr_tool,
    read_attached_file_tool,
    transcribe_audio_tool,
]

EXECUTOR_SYSTEM_PROMPT = (
    "You are the execution stage of a GAIA benchmark agent. Follow the given plan, "
    "using tools as needed to gather facts -- never guess a fact you can look up or compute. "
    "If a file is attached, you MUST inspect it with the appropriate tool before answering. "
    "End with a single, direct, concise final answer: no explanation, no extra commentary."
)

PLANNER_PROMPT = PromptTemplate(
    "You are the planning stage of a GAIA benchmark question-answering agent.\n"
    "Question: {question}\n"
    "{file_section}"
    "{history_section}"
    "{feedback_section}"
    "Available tools: you_answer, web_search, fetch_url_text, describe_image, extract_text_from_image, "
    "read_attached_file, transcribe_audio, calculate.\n"
    "Produce a concrete, ordered plan (as discrete steps) to answer the question."
)

# Runs once at the end, after the executor's raw answer has passed the judge --
# pulls out just the final answer, since GAIA grades exact string match and the
# executor's raw response is often wrapped in a sentence or two of explanation.
EXTRACT_ANSWER_PROMPT = PromptTemplate(
    "Question: {question}\n"
    "Response: {raw_answer}\n\n"
    "Extract ONLY the final answer to the question from the response above. "
    "No reasoning, no explanation, no leading/trailing words -- just the answer itself."
)


# Fired to (re-)request a plan; carries judge feedback when it's a retry, empty on the first attempt.
class PlanRequestEvent(Event):
    feedback: str | None = None


# Fired once the planner has produced a Plan, to hand off to the executor.
class PlanReadyEvent(Event):
    plan: Plan


# Fired once the executor has produced an answer, to hand off to the judge.
class ExecutionDoneEvent(Event):
    plan: Plan
    execution: ExecutionResult


# Fired when the judge accepts an answer, to hand off to the formatting step.
class FormatRequestEvent(Event):
    plan: Plan
    execution: ExecutionResult


# Turns past failed attempts into a short prompt section so the planner doesn't repeat them.
def _format_history(history: list[AttemptRecord]) -> str:
    if not history:
        return ""
    lines = ["Previous attempts on this question (avoid repeating these mistakes):"]
    for i, record in enumerate(history, start=1):
        lines.append(f"{i}. Answered {record.execution.raw_answer!r} -- judge said: {record.verdict.feedback}")
    return "\n".join(lines) + "\n"


# Builds the plain-text task description handed to the executor agent.
def _build_executor_task(question: str, plan: Plan, file_path: str | None) -> str:
    lines = [f"Question: {question}", "Plan:"]
    lines += [f"- {step_text}" for step_text in plan.steps]
    if file_path:
        lines.append(f"An attached file is available at this local path: {file_path}")
    return "\n".join(lines)


class GaiaWorkflow(Workflow):
    # Builds the tool-calling executor agent once, reused across every .run() call
    # (the plan/tools/LLM don't change between questions -- only Context state does).
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._executor_agent = FunctionAgent(
            tools=EXECUTOR_TOOLS,
            llm=get_llm(),
            system_prompt=EXECUTOR_SYSTEM_PROMPT,
            # streaming=False routes tool calls through achat_with_tools -> achat,
            # which ProviderChainLLM overrides. The default (streaming=True) instead
            # calls astream_chat, which is NOT overridden -- silently skipping the
            # whole fallback chain for every executor call. ToolCallResult events
            # (which execute_step reads below) come from a separate step regardless
            # of this flag, so nothing is lost by turning streaming off.
            streaming=False,
        )

    # Entry point: seeds Context with the question/file/retry-limit, then kicks off planning.
    @step
    async def start_step(self, ctx: Context[GaiaState], ev: StartEvent) -> PlanRequestEvent:
        async with ctx.store.edit_state() as state:
            state.question = ev.question
            state.file_path = getattr(ev, "file_path", None)
            state.max_retries = MAX_RETRIES
        print(f"\n[start] question={ev.question!r} file_path={getattr(ev, 'file_path', None)!r}")
        return PlanRequestEvent()

    # Asks the LLM to produce a structured Plan for the question (first attempt or a retry).
    @step
    async def plan_step(self, ctx: Context[GaiaState], ev: PlanRequestEvent) -> PlanReadyEvent:
        state = await ctx.store.get_state()
        file_section = f"An attachment is available at: {state.file_path}\n" if state.file_path else ""
        feedback_section = f"Feedback to address from the previous attempt: {ev.feedback}\n" if ev.feedback else ""
        # astructured_predict asks the LLM to return data shaped exactly like the
        # Plan model -- no manual JSON parsing, the LLM's output is validated into a Plan directly.
        plan = await get_llm().astructured_predict(
            Plan,
            PLANNER_PROMPT,
            question=state.question,
            file_section=file_section,
            history_section=_format_history(state.history),
            feedback_section=feedback_section,
        )
        print(f"[plan] {len(plan.steps)} step(s): {plan.steps}")
        return PlanReadyEvent(plan=plan)

    # Runs the tool-calling executor agent on the plan and captures its answer + tool-call trace.
    @step
    async def execute_step(self, ctx: Context[GaiaState], ev: PlanReadyEvent) -> ExecutionDoneEvent:
        state = await ctx.store.get_state()
        task = _build_executor_task(state.question, ev.plan, state.file_path)

        # agent.run() returns a handler immediately; streaming its events lets us watch
        # each tool call as it happens, then awaiting the handler gives the final answer.
        handler = self._executor_agent.run(user_msg=task)
        trace_lines: list[str] = []
        tool_names: list[str] = []
        async for event in handler.stream_events():
            if isinstance(event, ToolCallResult):
                print(f"[execute] tool call: {event.tool_name}({event.tool_kwargs}) -> {str(event.tool_output)[:200]}")
                tool_names.append(event.tool_name)
                trace_lines.append(f"Called {event.tool_name}({event.tool_kwargs}) -> {str(event.tool_output)[:300]}")
        response = await handler
        print(f"[execute] raw answer: {str(response)!r}")

        execution = ExecutionResult(
            raw_answer=str(response),
            reasoning_trace="\n".join(trace_lines) if trace_lines else "No tools were called; answered directly.",
            tool_calls_made=tool_names,
        )
        return ExecutionDoneEvent(plan=ev.plan, execution=execution)

    # Judges the executor's answer; loops back to planning with feedback if it fails and retries remain.
    @step
    async def eval_step(self, ctx: Context[GaiaState], ev: ExecutionDoneEvent) -> PlanRequestEvent | FormatRequestEvent:
        state = await ctx.store.get_state()
        verdict = judge_answer(
            question=state.question,
            plan_summary="; ".join(ev.plan.steps),
            reasoning_trace=ev.execution.reasoning_trace,
            candidate_answer=ev.execution.raw_answer,
        )
        print(f"[eval] is_valid={verdict.is_valid} confidence={verdict.confidence:.2f} feedback={verdict.feedback!r}")

        # edit_state() is a transactional block: mutations only become visible to other
        # steps once the block exits, so retry_count/history stay consistent under the loop.
        async with ctx.store.edit_state() as edit:
            edit.retry_count += 1
            edit.history.append(AttemptRecord(plan=ev.plan, execution=ev.execution, verdict=verdict))
            retry_count, max_retries = edit.retry_count, edit.max_retries

        if not verdict.is_valid and retry_count < max_retries:
            print(f"[eval] retrying ({retry_count}/{max_retries})")
            return PlanRequestEvent(feedback=verdict.feedback)
        return FormatRequestEvent(plan=ev.plan, execution=ev.execution)

    # Extracts the concise final answer, applies GAIA formatting rules, sanity-checks
    # the result, and either stops or loops back.
    @step
    async def format_step(self, ctx: Context[GaiaState], ev: FormatRequestEvent) -> StopEvent | PlanRequestEvent:
        state = await ctx.store.get_state()

        # Strip the executor's raw response down to just the answer before the
        # mechanical formatting rules run -- GAIA grades exact match, so a full
        # sentence of explanation would never score even if factually correct.
        extracted = await get_llm().astructured_predict(
            ConciseAnswer,
            EXTRACT_ANSWER_PROMPT,
            question=state.question,
            raw_answer=ev.execution.raw_answer,
        )
        formatted = format_answer(extracted.answer)
        print(f"[format] extracted={extracted.answer!r} -> formatted={formatted!r}")

        # Reuses the same judge as a final sanity check on the *formatted* string, in case
        # extraction or mechanical cleanup accidentally changed its meaning.
        sanity = judge_answer(
            question=state.question,
            plan_summary="; ".join(ev.plan.steps),
            reasoning_trace=ev.execution.reasoning_trace,
            candidate_answer=formatted,
        )
        print(f"[format] sanity check is_valid={sanity.is_valid} feedback={sanity.feedback!r}")
        if sanity.is_valid:
            return StopEvent(result=formatted)

        async with ctx.store.edit_state() as edit:
            edit.retry_count += 1
            retry_count, max_retries = edit.retry_count, edit.max_retries

        if retry_count < max_retries:
            print(f"[format] retrying ({retry_count}/{max_retries})")
            return PlanRequestEvent(feedback=sanity.feedback)

        # Retries exhausted -- submit the best formatted answer rather than nothing.
        return StopEvent(result=formatted)