import json
import traceback
from pathlib import Path

import click
import pandas as pd

import minimal.data as datasets
from minimal.configuration import cfg
from minimal.evaluation import eval_dataset
from minimal.instrumentation.arize import instrument_arize
from minimal.searchspace import SearchSpace
from minimal.tuner import build_flow


def determine_output_filepath(
    output_file_option: Path | None, dataset_name: str
) -> Path:
    """Determines the final output CSV file path."""
    if output_file_option is None:
        output_file = Path(f"./flowgen-eval-results-{dataset_name}.csv")
    else:
        output_file = output_file_option

    output_file.parent.mkdir(parents=True, exist_ok=True)

    click.secho(f"Output file set to: {output_file}", fg="blue")
    return output_file


def load_existing_results(output_file: Path) -> tuple[pd.DataFrame, set[str]]:
    """Loads existing results from the CSV file, returns DataFrame and set of processed flow strings."""
    results_df = pd.DataFrame()
    processed_flow_strs = set()

    if not output_file.exists():
        return results_df, processed_flow_strs

    results_df = pd.read_csv(output_file)
    if results_df.empty:
        return (results_df, processed_flow_strs)

    if "flow_params_str" in results_df.columns:
        # Ensure correct type and handle potential NaN values before adding to set
        results_df["flow_params_str"] = results_df["flow_params_str"].astype(str)
        processed_flow_strs = set(results_df["flow_params_str"].dropna().unique())

    return results_df, processed_flow_strs


def load_flow_definitions(flows_file: Path) -> list[dict]:
    """Loads flow definitions from a JSON-lines file."""
    with open(flows_file) as flows_data:
        flow_jsons = flows_data.readlines()
    flow_defs = [json.loads(data) for data in flow_jsons]
    click.secho(
        f"Loaded {len(flow_defs)} flow definitions from {flows_file}", fg="green"
    )
    return flow_defs


def initialize_instrumentation():
    """Initializes Arize instrumentation if available."""
    try:
        instrument_arize()
        click.secho(
            f"Instrumentation enabled. Run `phoenix serve` in another shell if you haven't, "
            f"and find results at {cfg.instrumentation.arize_endpoint.rstrip('v1/traces')}",
            fg="cyan",
        )
    except ImportError:
        click.secho(
            "Instrumentation library (e.g., Arize) not found or import failed. Skipping.",
            fg="yellow",
        )


@click.command()
@click.option(
    "-f",
    "--flows-file",
    type=click.Path(exists=True, dir_okay=False, path_type=Path),
    required=True,
    help="Path to the file containing flow definitions (one JSON per line).",
)
@click.option(
    "-d",
    "--dataset-name",
    type=str,
    default="HotPotQAHF",
    show_default=True,
    help="Name of the dataset class in minimal.data.",
)
@click.option(
    "-i",
    "--instrumentation",
    is_flag=True,
    default=False,
    help="Enable Arize instrumentation.",
)
@click.option(
    "-o",
    "--output-file",
    default=None,
    type=click.Path(dir_okay=False, path_type=Path),
    help="Path to the output CSV file. Defaults to ./flowgen-eval-results-{dataset_name}.csv",
)
@click.option(
    "-m",
    "--max-evals",
    type=int,
    default=1000,
    show_default=True,
    help="Maximum QA pairs to evaluate per flow",
)
def main(
    flows_file: Path,
    dataset_name: str,
    instrumentation: bool,
    output_file: Path | None,
    max_evals: int,
):
    """
    Evaluates flows defined in flows_file against the specified dataset,
    stores results in a CSV, appends results iteratively, and skips
    flows already present in the output CSV.
    """
    if instrumentation:
        initialize_instrumentation()

    # --- Setup ---
    output_file = determine_output_filepath(output_file, dataset_name)
    results_df, processed_flow_strs = load_existing_results(output_file)
    flow_defs = load_flow_definitions(flows_file)

    try:
        dataset_class = getattr(datasets, dataset_name)
    except AttributeError:
        click.secho(
            f"Error: Dataset class '{dataset_name}' not found in minimal.data.",
            fg="red",
            err=True,
        )
        raise
    dataset = dataset_class()

    total_flows = len(flow_defs)
    processed_count = 0
    skipped_count = 0
    error_count = 0

    for i, flow_params in enumerate(flow_defs):
        flow_params_str = ""
        try:
            flow_params_str = json.dumps(flow_params, sort_keys=True)

            if flow_params_str in processed_flow_strs:
                skipped_count += 1
                click.secho(
                    f"({i + 1}/{total_flows}) Skipping already processed flow: {flow_params_str[:80]}...",
                    fg="cyan",
                )
                continue

            click.secho(
                f"({i + 1}/{total_flows}) Processing flow: {flow_params_str[:80]}...",
                fg="yellow",
            )

            search_space = SearchSpace()

            click.secho("  Building flow...", fg="white", nl=False)
            flow = build_flow(flow_params, search_space, dataset)
            click.secho("  Done.", fg="green")

            click.secho(
                f"  Evaluating dataset '{dataset_name}'...", fg="white", nl=False
            )
            results = eval_dataset(dataset, flow, max_evals=max_evals)
            click.secho("  Done.", fg="green")
            click.echo(f"  Results: {json.dumps(results, sort_keys=True, indent=4)}")

            new_row_data = {"flow_params_str": flow_params_str}
            new_row_data.update(results)

            new_row_df = pd.DataFrame([new_row_data])
            results_df = pd.concat([results_df, new_row_df], ignore_index=True)

            results_df.to_csv(output_file, index=False)
            click.secho(f"  Results appended and saved to {output_file}", fg="blue")

            processed_flow_strs.add(flow_params_str)
            processed_count += 1

        except Exception as e:
            error_count += 1
            click.secho(
                f"\nError processing flow ({flow_params_str[:80]}...): {e}",
                fg="red",
                err=True,
            )
            click.secho(traceback.format_exc(), fg="red", err=True)

    click.secho("\n--- Processing Summary ---", fg="white", bold=True)
    click.secho(f"Total flows in file:      {total_flows}", fg="white")
    click.secho(f"Successfully processed:   {processed_count}", fg="green")
    if skipped_count:
        click.secho(f"Skipped (already done): {skipped_count}", fg="cyan")
    if error_count:
        click.secho(
            f"Errors encountered:       {error_count}",
            fg="red",
        )
    click.secho(f"{results_df.describe()}", fg="blue")
    click.secho(f"Final results saved to:   {output_file}", fg="blue")


if __name__ == "__main__":
    main()
