"""Command-line interface for running SBC analysis."""
import pathlib
import tomllib

import click
import cmdstanpy

from . import stan


def parse_toml_config(toml_path):
    """Parse TOML configuration file and extract SBC parameters.

    Parameters
    ----------
    toml_path : str or pathlib.Path
        Path to the TOML configuration file.

    Returns
    -------
    config : dict
        Dictionary containing all SBC configuration parameters.
    """
    toml_path = pathlib.Path(toml_path)

    with open(toml_path, "rb") as f:
        config = tomllib.load(f)

    return config


@click.command()
@click.argument("config_file", type=click.Path(exists=True))
def main(config_file):
    """Run SBC analysis from a TOML configuration file.

    All parameters including output file path are specified in the config file.

    CONFIG_FILE: Path to TOML configuration file
    """
    click.echo(f"Reading configuration from: {config_file}")
    config = parse_toml_config(config_file)

    # Get required fields from TOML file
    required_keys = (
        'prior_predictive_stan_file',
        'posterior_stan_file',
        'prior_predictive_model_data',
        'posterior_model_data',
        'measured_data'
    )
    sbc_kwargs = {}
    for key in required_keys:
        if key in config:
            sbc_kwargs[key] = config[key]
        else:
            raise RuntimeError(f'{key} must be specified.')

    # Get optional fields
    sbc_kwargs['var_names'] = config.get('var_names', None)
    sbc_kwargs['measured_data_dtypes'] = config.get('measured_data_dtypes', {})
    sbc_kwargs['posterior_predictive_var_names'] = config.get('posterior_predictive_var_names', None)
    sbc_kwargs['log_likelihood_var_name'] = config.get('log_likelihood_var_name', None)
    sbc_kwargs['sampling_kwargs'] = config.get('sampling_kwargs', {})
    sbc_kwargs['diagnostic_check_kwargs'] = config.get('diagnostic_check_kwargs', {})
    sbc_kwargs['cores'] = config.get('cores', 1)
    sbc_kwargs['N'] = config.get('N', 400)
    sbc_kwargs['n_prior_draws_for_sd'] = config.get('n_prior_draws_for_sd', 1000)
    sbc_kwargs['samples_dir'] = config.get('samples_dir', 'sbc_samples')
    sbc_kwargs['remove_sample_files'] = config.get('remove_sample_files', True)
    sbc_kwargs['progress_bar'] = config.get('progress_bar', False)
    sbc_kwargs['output_csv'] = config.get('output_csv', 'sbc_results.csv')
    sbc_kwargs['overwrite'] = config.get('overwrite', False)

    # Run SBC (results written directly to CSV as workers complete)
    click.echo("Running SBC analysis...")
    click.echo(f"Results will be written to: {sbc_kwargs['output_csv']}\n")

    result_df = stan.sbc(**sbc_kwargs)

    click.echo(f"\nSBC analysis complete! Results saved to: {sbc_kwargs['output_csv']}")
    click.echo(f"Total rows: {len(result_df)}")


if __name__ == "__main__":
    main()
