import csv
import random
import statistics
import sys
import time
from collections import defaultdict
from functools import wraps

import typer
from icat_plus_client.models.session import Session

from data import dataset, investigation, sample, session

performance_app = typer.Typer(help="Performance workbench commands")


def get_role(user: Session):
    if user.is_administrator:
        return "administrator"
    if user.is_instrument_scientist:
        return "instrumentScientist"
    return "user"


# Decorator to measure execution time
def timed(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        start_time = time.time()
        result = func(*args, **kwargs)
        end_time = time.time()
        duration = (end_time - start_time) * 1000
        return result, round(duration)

    return wrapper


@timed
def get_timed_datasets(**kwargs):
    return dataset.get_datasets(**kwargs)


@timed
def get_timed_samples(**kwargs):
    return sample.get_samples_by(**kwargs)


def get_timed_dataset_by_samples(**kwargs):
    return sample.get_samples_by(**kwargs)


def workbench(
    fetch_fn,
    result_label: str,
    *,
    random_seed: int = 42,
    cooldown_seconds: float = 0,
    skip_range=(0, 100),
    limit_range=(100, 200),
    x: int = 1,
    y: int = 1,
):
    """
    Generic benchmark decorator.

    fetch_fn: timed callable (get_timed_datasets / get_timed_samples)
    result_label: CSV column name ('datasets' / 'samples')
    """

    def decorator(func):
        @wraps(func)
        def wrapper(
            tokens: str = typer.Option(..., "-t", "--token"),
            investigation_ids: str | None = typer.Option(None, "-i", "--investigation-ids"),
            start_date: str = typer.Option(
                None,
                "-s",
                "--start-date",
                help="Start date range. Format is YYYY-MM-DD. Example: 2019-01-01",
            ),
            end_date: str = typer.Option(
                None,
                "-e",
                "--end-date",
                help="End date range. Format is YYYY-MM-DD. Example: 2019-01-01",
            ),
            skip_count: int | None = typer.Option(
                1,
                "-k",
                "--skip-count",
                help="Number of skip variations",
            ),
            limit_count: int | None = typer.Option(
                1,
                "-l",
                "--limit-count",
                help="Number of limit variations",
            ),
            output: str | None = typer.Option(None, "-o", "--output"),
        ):
            user_tokens = tokens.split(",")

            investigation_id_list = []
            if investigation_ids is not None:
                investigation_id_list = investigation_ids.split(",")
            else:
                investigations = investigation.get_investigation_by_id(
                    user_tokens[0], start_date=start_date, end_date=end_date
                )
                investigation_id_list = [i.id for i in investigations]

            fields = [
                "role",
                "user",
                "investigation_id",
                "skip",
                "limit",
                result_label,
                "time(ms)",
            ]

            stdout_writer = csv.DictWriter(sys.stdout, fieldnames=fields, delimiter="\t")
            stdout_writer.writeheader()

            file_writer = None
            f = None
            if output:
                f = open(output, "w", newline="")
                file_writer = csv.DictWriter(f, fieldnames=fields, delimiter="\t")
                file_writer.writeheader()

            role_timings = defaultdict(list)

            try:
                random.seed(random_seed)

                # --------------------------------------------------------------
                # Build randomized calls
                # --------------------------------------------------------------
                calls = []
                for investigation_id in investigation_id_list:
                    for _ in range(skip_count):
                        skip = str(random.randint(*skip_range))
                        for _ in range(limit_count):
                            limit = str(random.randint(*limit_range))
                            for token in user_tokens:
                                calls.append(
                                    {
                                        "token": token,
                                        "investigation_id": investigation_id,
                                        "skip": skip,
                                        "limit": limit,
                                    }
                                )

                random.shuffle(calls)

                # --------------------------------------------------------------
                # Execute benchmark
                # --------------------------------------------------------------
                for call in calls:
                    token = call["token"]
                    investigation_id = call["investigation_id"]
                    skip = call["skip"]
                    limit = call["limit"]

                    user_session: Session = session.get_info(token)
                    role = get_role(user_session.user)

                    result, duration = fetch_fn(
                        token=token,
                        investigation_ids=investigation_id,
                        limit=limit,
                        skip=skip,
                    )

                    role_timings[role].append(duration)

                    row = {
                        "role": role,
                        "user": user_session.user.username,
                        "investigation_id": investigation_id,
                        "skip": skip,
                        "limit": limit,
                        result_label: len(result),
                        "time(ms)": duration,
                    }

                    stdout_writer.writerow(row)
                    if file_writer:
                        file_writer.writerow(row)

                    time.sleep(cooldown_seconds)

            finally:
                if f:
                    f.close()

            # --------------------------------------------------------------
            # Statistics
            # --------------------------------------------------------------
            print("\n=== Benchmark statistics per role ===\n")
            for role, times in role_timings.items():
                if len(times) < 2:
                    continue
                print(
                    f"Role: {role} | "
                    f"count={len(times)} | "
                    f"avg={statistics.mean(times):.2f}ms | "
                    f"min={min(times):.2f}ms | "
                    f"max={max(times):.2f}ms | "
                    f"std={statistics.stdev(times):.2f}ms"
                )

        return wrapper

    return decorator


@performance_app.command("dataset")
@workbench(
    fetch_fn=get_timed_datasets,
    result_label="datasets",
)
def dataset_workbench(
    tokens: str = typer.Option(..., "-t", "--token", help="Comma separated list of ICAT tokens"),
    investigation_ids: str | None = typer.Option(
        None,
        "--investigation-ids",
        "-i",
        help="Comma separated list of investigation IDs (optional)",
    ),
    start_date: str = typer.Option(
        None,
        "-s",
        "--start-date",
        help="Start date range. Format is YYYY-MM-DD. Example: 2019-01-01",
    ),
    end_date: str = typer.Option(
        None,
        "-e",
        "--end-date",
        help="End date range. Format is YYYY-MM-DD. Example: 2019-01-01",
    ),
    skip_count: int | None = typer.Option(
        1,
        "-k",
        "--skip-count",
        help="Number of skip variations",
    ),
    limit_count: int | None = typer.Option(
        1,
        "-l",
        "--limit-count",
        help="Number of limit variations",
    ),
    output: str | None = typer.Option(
        None,
        "-o",
        "--output",
        help="CSV output file path. Example: myfile.csv",
    ),
):
    """Benchmark dataset retrieval."""


@performance_app.command(
    "sample",
    help="""
    Generates volume statistics for all instruments over a date range.

    Example usage:

    .. code-block:: bash

      minedrac performance sample --token <TOKEN>  --start-date 2019-02-01 --end-date 2019-02-02

    .. code-block:: bash

      === Benchmark statistics per role ===

        Role: instrumentScientist | count=28 | avg=155.68ms | min=62.00ms | max=856.00ms | std=217.77ms
        Role: administrator | count=28 | avg=169.14ms | min=64.00ms | max=583.00ms | std=180.10ms
        Role: user | count=28 | avg=117.39ms | min=66.00ms | max=560.00ms | std=125.73ms
                               """,
)
@workbench(
    fetch_fn=get_timed_samples,
    result_label="samples",
)
def samples_workbench(
    tokens: str = typer.Option(..., "-t", "--token", help="Comma separated list of ICAT tokens"),
    investigation_ids: str | None = typer.Option(
        None,
        "--investigation-ids",
        "-i",
        help="Comma separated list of investigation IDs (optional)",
    ),
    start_date: str = typer.Option(
        None,
        "-s",
        "--start-date",
        help="Start date range. Format is YYYY-MM-DD. Example: 2019-01-01",
    ),
    end_date: str = typer.Option(
        None,
        "-e",
        "--end-date",
        help="End date range. Format is YYYY-MM-DD. Example: 2019-01-01",
    ),
    skip_count: int | None = typer.Option(
        1,
        "-k",
        "--skip-count",
        help="Number of skip variations",
    ),
    limit_count: int | None = typer.Option(
        1,
        "-l",
        "--limit-count",
        help="Number of limit variations",
    ),
    output: str | None = typer.Option(
        None,
        "-o",
        "--output",
        help="CSV output file path. Example: myfile.csv",
    ),
):
    """Benchmark dataset retrieval."""


# For click documentation
click_performance_app = typer.main.get_command(performance_app)
