from pathlib import Path

from flock.core import Flock, FlockFactory
from flock.core.logging.logging import get_logger
from flock.core.tools.basic_tools import (
    code_eval,
    evaluate_math,
    get_web_content_as_markdown,
    web_search_bing,
    web_search_duckduckgo,
    web_search_tavily,
)

logger = get_logger("evaluate_example")

# --- 1. Define the Agent to Evaluate ---
# Let's create a simple Q&A agent that takes a 'query' and outputs an 'answer'.
# We'll map the dataset's 'question' column to our agent's 'query' input.
qa_agent = FlockFactory.create_default_agent(
    name="qa_agent",
    description="Answers questions based on general knowledge.",
    input="query: str | The question to be answered",
    output="answer: str | The concise answer to the question",
    tools=[web_search_tavily,
           web_search_duckduckgo,
           web_search_bing,
           get_web_content_as_markdown,
           code_eval,evaluate_math],
    model="openai/gpt-4o", # Or your preferred model
    max_tokens=15000
)

# --- 2. (Optional) Define an LLM Judge Agent ---
# This agent will compare the predicted vs. actual answer based on instructions.
judge_agent = FlockFactory.create_default_agent(
    name="answer_judge",
    description="Evaluates the correctness of a predicted answer compared to a reference answer.",
    input="""
        prediction: str | The answer generated by the agent being evaluated,
        reference: str | The ground truth answer from the dataset,
        question: str | The question that was asked,
        instruction: str | Optional instructions for evaluation (e.g., 'Focus on factual accuracy')
    """,
    output="""
        score: float | A score between 0.0 (incorrect) and 1.0 (perfectly correct),
        reasoning: str | Explanation for the score given
    """,
    model="openai/gpt-4o", # Use a capable model for judging
    max_tokens=15000
)

# --- 3. Create the Flock Instance ---
flock = Flock(name="evaluation_flock", enable_logging=False)
flock.add_agent(qa_agent)
flock.add_agent(judge_agent) # Add the judge agent if using it as a metric

# --- 4. Define Evaluation Parameters ---

# Dataset Identifier (Hugging Face Hub)
# https://huggingface.co/datasets/smolagents/benchmark-v1
DATASET_ID = "smolagents/benchmark-v1"

# Input Mapping: Map dataset 'question' column to agent's 'query' input
input_mapping = {
    "question": "query"
}

# Answer Mapping: Map agent's 'answer' output to dataset's 'true_answer' column
answer_mapping = {
    "answer": "true_answer"
}

# Metrics to calculate: Use built-in names and the judge agent instance
# Note: For 'semantic_similarity', ensure 'sentence-transformers' is installed.
# Note: For 'fuzzy_match', ensure 'thefuzz[speedup]' is installed.
# Note: For 'rouge_*', ensure 'rouge-score' is installed.
metrics_to_run = [
    "exact_match",
    "fuzzy_match",
    "semantic_similarity",
    judge_agent # Pass the actual agent instance for LLM judging
]

# Optional: Configuration for specific metrics (e.g., threshold for fuzzy match)
metric_configs = {
    "fuzzy_match": {"threshold": 90},
    # 'llm_judge': {'instruction': 'Assess if the prediction correctly answers the core intent of the reference.'} # Config for the judge agent itself, or pass via static_inputs if needed by the judge's prompt
}

# Optional: Columns from the dataset to include in the results file
metadata_columns_to_keep = ["source", "true_reasoning"]

# Optional: Output file path
results_file = Path("evaluation_results.csv")

# --- 5. Run the Evaluation ---
print(f"Starting evaluation of agent '{qa_agent.name}' on dataset '{DATASET_ID}'...")

# Using the synchronous wrapper for simplicity in this example
try:
    evaluation_df = flock.evaluate(
        dataset=DATASET_ID,                       # HF Dataset ID
        # dataset_split=DATASET_SPLIT,            # TODO: Add support for dataset splits
        start_agent=qa_agent.name,                # Agent to run
        input_mapping=input_mapping,              # How inputs map
        answer_mapping=answer_mapping,            # How outputs map to answers
        metrics=metrics_to_run,                   # What metrics to compute
        metric_configs=metric_configs,            # Metric-specific settings
        metadata_columns=metadata_columns_to_keep,# Keep 'source' column
        output_file=results_file,                 # Save detailed results
        return_dataframe=True,                    # Get results as DataFrame
        silent_mode=True,                         # Show progress bar
        error_handling='log'                      # Log errors but continue
    )

    print(f"\nEvaluation complete. Results saved to {results_file}")

    # --- 6. Display Summary ---
    print("\n--- Evaluation Summary ---")
    print(f"Total items evaluated: {len(evaluation_df)}")

    # Display average scores for numeric metrics
    numeric_metrics = evaluation_df.select_dtypes(include='number').columns
    # Filter to include only the calculated metric columns (based on metrics_to_run)
    metric_cols = [m if isinstance(m, str) else m.name for m in metrics_to_run]
    # Expand potential dict metrics (like rouge)
    flat_metric_names = []
    for m_name in metric_cols:
        matching_cols = [c for c in numeric_metrics if c.startswith(m_name)]
        if matching_cols:
            flat_metric_names.extend(matching_cols)
        elif m_name in numeric_metrics: # Handle simple metric names directly
             flat_metric_names.append(m_name)

    if flat_metric_names:
        avg_scores = evaluation_df[flat_metric_names].mean()
        print("\nAverage Scores:")
        print(avg_scores.to_string())
    else:
        print("No numeric metrics found to average.")

    # Display head of the detailed results DataFrame
    print("\n--- Detailed Results (First 5 Rows) ---")
    print(evaluation_df.head().to_markdown(index=False))

except ValueError as ve:
    print(f"\n[ERROR] Evaluation setup failed: {ve}")
except Exception as e:
    print(f"\n[ERROR] An unexpected error occurred during evaluation: {e}")