# -*- coding: utf-8 -*-
"""Example of training an Email Search agent with Trinity-RFT."""
import os
from typing import Dict
from _email_search_agent import EmailSearchAgent
from _utils import (  # pylint: disable=E0611
    AnswerModel,
    FinalRubric,
    QueryModel,
)
from agentscope import logger
from agentscope.formatter import OpenAIChatFormatter
from agentscope.message import Msg
from agentscope.tuner import (
    TunerModelConfig,
    DatasetConfig,
    JudgeOutput,
    WorkflowOutput,
    AlgorithmConfig,
    tune,
)
from agentscope.model import ChatModelBase


SYSTEM_PROMPT = """You are an email search agent. You are given a user query
and a list of tools you can use to search the user's email. Use the tools to
search the user's emails and find the answer to the user's query. You may take
up to {max_turns} turns to find the answer, so if your first seach doesn't
find the answer, you can try with different keywords.

Always describe what you see and plan your next steps clearly. When taking
actions, explain what you're doing and why. When the answer to the task is
found, call `generate_response` to finish the process. Only call
`generate_response` when answer is found. You should not respond any next steps
in `generate_response`. Complete all steps and then call `generate_response`.

User's email address is {inbox_address}

Today's date is {query_date}

"""


async def run_email_search_agent(
    task: Dict,
    model: ChatModelBase,
    auxiliary_models: Dict[str, ChatModelBase],
) -> WorkflowOutput:  # noqa: PLR0915
    """A workflow function using the Email Search agent to solve tasks.

    Args:
        task (Dict): The task to be solved.
            Should contain fields from QueryModel.
        model (TrinityChatModel): The language model to use.

    Returns:
        WorkflowOutput: The output containing the agent's response.
    """
    assert len(auxiliary_models) > 0, "LLM-as-a-Judge is required"

    # Parse task data
    query = QueryModel.model_validate(task)
    question = task.get("question", task.get("task_desc", ""))

    # Get workflow arguments with defaults
    workflow_args = task.get("workflow_args", {})
    max_turns = int(workflow_args.get("max_turns", 10))

    # Format system prompt
    system_prompt = SYSTEM_PROMPT.format(
        max_turns=max_turns,
        inbox_address=query.inbox_address,
        query_date=query.query_date,
    )

    # Create EmailSearchAgent
    agent = EmailSearchAgent(
        name="email_search_agent",
        sys_prompt=system_prompt,
        model=model,
        formatter=OpenAIChatFormatter(),
        max_iters=max_turns,
    )

    # Reset agent state for a new rollout
    await agent.reset()

    # Run the agent with structured output
    response = await agent.reply(
        msg=Msg("user", question, role="user"),
        structured_model=AnswerModel,
    )

    # Extract answer and sources from response metadata
    answer_and_sources = response.metadata or {}
    if not answer_and_sources:
        # Fallback: try to parse from content
        answer_and_sources = {
            "answer": response.get_text_content() or "",
            "sources": [],
        }

    # Store agent state for judge function
    # We'll pass this through the response metadata
    response_metadata = {
        "answer_and_sources": answer_and_sources,
        "query": query.model_dump(),
        "message_id_list": agent.message_id_list,
        "ever_read_message_ids": agent.ever_read_message_ids,
        # Estimate actual_turns from memory length
        "actual_turns": (
            max(1, (len(agent.memory.content) - 1) // 2)
            if len(agent.memory.content) > 1
            else 1
        ),
    }

    # Update response metadata
    if response.metadata is None:
        response.metadata = {}
    response.metadata.update(response_metadata)

    return WorkflowOutput(
        response=response,
    )


def _calculate_partial_rewards(rubric: FinalRubric) -> float:
    """Calculate partial rewards based on rubric."""
    partial_rewards = 0.0
    partial_rewards += 0.1 if rubric.ever_found_right_email else 0
    partial_rewards += 0.1 if rubric.ever_read_right_email else 0
    partial_rewards += 0.1 if rubric.sources_correct else 0
    return partial_rewards


def _calculate_correct_answer_reward(
    rubric: FinalRubric,
    max_turns: int,
) -> float:
    """Calculate reward for correct answers."""
    reward = 1.0
    reward += 0.3 if rubric.sources_correct else 0
    reward += 0.1 / rubric.num_sources if rubric.num_sources > 0 else 0
    reward += 0.1 * (1 - rubric.num_turns / max_turns)
    return reward


def _initialize_rubric(
    answer: str,
    sources: list[str],
    actual_turns: int,
    query: QueryModel,
    message_id_list: list[str],
    ever_read_message_ids: list[str],
) -> FinalRubric:
    """Initialize and populate rubric with basic information."""
    rubric = FinalRubric()
    rubric.attempted_answer = answer is not None and answer != ""
    rubric.returned_i_dont_know = answer == "I don't know"
    rubric.num_sources = len(sources)
    rubric.num_turns = actual_turns

    if len(query.message_ids) > 0:
        rubric.ever_found_right_email = query.message_ids[0] in message_id_list
        rubric.ever_read_right_email = (
            query.message_ids[0] in ever_read_message_ids
        )
        rubric.sources_correct = query.message_ids[0] in sources
    return rubric


async def email_search_judge(
    task: Dict,
    response: Msg,
    auxiliary_models: Dict[str, ChatModelBase],
) -> JudgeOutput:
    """A judge function to calculate reward based on agent's response.

    Args:
        task (Dict): The task information for the corresponding workflow.
        response (Msg): The response generated by the corresponding workflow.
        auxiliary_models (Dict[str, ChatModelBase]):
            A dictionary of additional chat models available for LLM-as-a-Judge
            usage. The keys are model names, and the values are the
            corresponding ChatModelBase instances.

    Returns:
        JudgeOutput: The reward value assigned by the judge function.
    """
    # Extract metadata from response
    metadata = response.metadata or {}
    answer_and_sources = metadata.get("answer_and_sources", {})
    query_dict = metadata.get("query", {})
    message_id_list = metadata.get("message_id_list", [])
    ever_read_message_ids = metadata.get("ever_read_message_ids", [])
    actual_turns = metadata.get("actual_turns", 0)

    # Parse query model
    if not query_dict:
        query_dict = task
    query = QueryModel.model_validate(query_dict)

    # Get arguments
    workflow_args = task.get("workflow_args", {})
    max_turns = int(workflow_args.get("max_turns", 10))

    # Extract answer and sources
    try:
        answer = answer_and_sources.get("answer", None)
        sources = answer_and_sources.get("sources", [])
    except Exception:
        result = {"accuracy": 0.0, "format": -1.0}
        return JudgeOutput(
            reward=sum(result.values()),
            metrics=result,
        )

    if answer is None:
        result = {"accuracy": 0.0, "format": -1.0}
        return JudgeOutput(
            reward=sum(result.values()),
            metrics=result,
        )

    # Initialize rubric
    rubric = _initialize_rubric(
        answer,
        sources,
        actual_turns,
        query,
        message_id_list,
        ever_read_message_ids,
    )

    # Judge correctness using LLM-as-a-Judge
    try:
        judge_model = (
            auxiliary_models.get("judge") or list(auxiliary_models.values())[0]
            if auxiliary_models
            else None
        )
        judge_response = await judge_correctness(
            answer,
            query,
            judge_model,
        )
        rubric.answer_correct = judge_response
    except Exception as e:
        logger.error("Error judging correctness: %s", e)
        rubric.answer_correct = False

    # Calculate rewards
    partial_rewards = _calculate_partial_rewards(rubric)

    if rubric.attempted_answer and not rubric.answer_correct:
        result = {"accuracy": -1.0, "format": partial_rewards}
    elif rubric.returned_i_dont_know or rubric.ran_out_of_turns:
        result = {"accuracy": 0.0, "format": partial_rewards}
    elif rubric.answer_correct:
        reward = _calculate_correct_answer_reward(rubric, max_turns)
        result = {"accuracy": 1.0, "format": reward}
    else:
        result = {"accuracy": 0.0, "format": 0.0}

    metrics = result.copy()
    metrics.update({"actual_turns": actual_turns})

    return JudgeOutput(
        reward=sum(result.values()),
        metrics=metrics,
    )


# LLM-as-a-judge


async def judge_correctness(
    answer: str,
    query: QueryModel,
    judge: ChatModelBase,
) -> bool:
    """Use an LLM to decide whether *answer* matches *query.answer*.

    Returns a boolean *accept* flag used for scoring.
    """

    system_prompt = """You are given a question, the reference answer
(labelled **Reference answer**), and an answer generated by an AI assistant
(labelled **AI answer**).

Follow these steps to decide whether the AI answer should be accepted:
1. Identify EXACTLY what information the **question** is asking for
   (e.g. who, what, when, where, why, how, quantity, etc.).
2. From the **Reference answer**, extract ONLY the facts that are required
   to directly satisfy the information need identified in step 1. Treat all
   other facts as non-essential context.
3. Verify that every essential fact from step 2 appears in the **AI answer**
   with the same meaning. Differences in wording, order, or additional
   non-conflicting details are allowed.
4. If any essential fact is missing or contradicted in the **AI answer**,
   then *accept* must be **false**. Otherwise *accept* must be **true**.

Important: Do NOT penalise the **AI answer** for omitting non-essential
facts that appear in the **Reference answer**. The answer should only be
rejected for errors or omissions in the information explicitly requested by
the question.

Return your judgement **accept** from **true** and **false**. Do not return
any other text or formatting.
"""
    prompt = (
        f"Question: {query.question}\n"
        f"Reference answer: {query.answer}\n"
        f"AI answer: {answer}"
    )

    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": prompt},
    ]
    chat_response = await judge(messages)

    # Extract text content from ChatResponse
    result_parts = []
    for block in chat_response.content:
        if isinstance(block, dict) and block.get("type") == "text":
            result_parts.append(str(block.get("text", "")))
    result = "".join(result_parts)
    logger.info("LLM judge response: %s", result)

    return "true" in result.lower()


# End of LLM-as-a-judge


if __name__ == "__main__":
    config_path = os.path.join(
        os.path.dirname(__file__),
        "config.yaml",
    )
    dataset = DatasetConfig(
        path="/path/to/enron_emails_dataset",
        split="train",
    )
    tuner_model = TunerModelConfig(
        model_path="Qwen/Qwen3-4B-Instruct-2507",
        max_model_len=20480,
        max_tokens=4096,
        inference_engine_num=4,
        reasoning_parser=None,
    )
    aux_models = {
        "judge": TunerModelConfig(
            model_path="Qwen/Qwen3-30B-A3B-Instruct-2507",
            max_model_len=2500,
            max_tokens=2048,
            inference_engine_num=1,
            tensor_parallel_size=2,
            tool_call_parser=None,
            reasoning_parser=None,
        ),
    }
    algorithm = AlgorithmConfig(
        algorithm_type="multi_step_grpo",
        group_size=8,
        batch_size=64,
        learning_rate=1e-6,
    )
    tune(
        workflow_func=run_email_search_agent,
        judge_func=email_search_judge,
        train_dataset=dataset,
        model=tuner_model,
        auxiliary_models=aux_models,
        algorithm=algorithm,
        config_path=config_path,
    )