import typing

import pandas
import pytablewriter

from cumulus_library import base_utils, databases, study_manifest
from cumulus_library.actions import builder
from cumulus_library.template_sql import base_templates


def run_generate_sql(
    config: base_utils.StudyConfig,
    manifest: study_manifest.StudyManifest,
    *,
    table_builder: str | None = None,
    db_parser: databases.DatabaseParser = None,
) -> None:
    """Generates reference SQL from table_builders listed in the manifest

    :param config: a StudyConfig object
    :param manifest: a StudyManifest object
    :keyword table_builder: a specific table_builder to target
    :keyword db_parser: a DB parser
    """
    all_generators = manifest.get_all_generators()
    doc_str = (
        "-- This sql was autogenerated as a reference example using the library\n"
        "-- CLI. Its format is tied to the specific database it was run against,\n"
        "-- and it may not be correct for all databases. Use the CLI's build \n"
        "-- option to derive the best SQL for your dataset."
    )
    for run_stage, file in enumerate(all_generators):
        if table_builder and file.find(table_builder) == -1:
            continue  # pragma: no cover
        builder._run_builder(
            config=config,
            manifest=manifest,
            filename=file,
            db_parser=db_parser,
            write_reference_sql=True,
            doc_str=doc_str,
            data_path=None,
            run_stage=run_stage,
        )


def run_generate_markdown(
    config: base_utils.StudyConfig,
    manifest: study_manifest.StudyManifest,
    *,
    db_parser: databases.DatabaseParser = None,
) -> None:
    """Generates reference SQL from builders listed in the manifest

    :param config: a StudyConfig object
    :param manifest: a StudyManifest object
    """

    cursor = config.db.cursor()
    query = base_templates.get_show_tables(
        schema_name=config.schema, prefix=f"{manifest.get_study_prefix()}__"
    )
    tables = [x[0] for x in cursor.execute(query).fetchall()]
    query = base_templates.get_column_datatype_query(
        schema_name=config.schema, table_names=tables, include_table_names=True
    )
    study_df = pandas.DataFrame(
        cursor.execute(query).fetchall(), columns=["Column", "Type", "Table"]
    )
    with open(
        manifest._study_path / f"{manifest.get_study_prefix()}_generated.md",
        "w",
    ) as f:
        table_list = sorted(study_df["Table"].unique())
        count_tables = [t for t in table_list if "__count_" in t]
        base_tables = [t for t in table_list if "__count_" not in t]
        if len(count_tables) > 0:
            f.write(f"## {manifest.get_study_prefix()} count tables\n\n")
            for table in count_tables:
                _write_md_table(table, study_df, f)
        if len(base_tables) > 0:
            f.write(f"## {manifest.get_study_prefix()} base tables\n\n")
            for table in base_tables:
                _write_md_table(table, study_df, f)


def _write_md_table(name: str, df: pandas.DataFrame, file: typing.IO):
    table_df = df[df["Table"] == name].drop("Table", axis=1)
    table_df = table_df.assign(Description="")
    writer = pytablewriter.MarkdownTableWriter(dataframe=table_df)
    writer.table_name = f"{name}\n"
    writer.set_indent_level(2)
    writer.stream = file
    writer.write_table()
    file.write("\n\n")
