import re, os
from pathlib import Path
from fastcore.script import Param, call_parse, store_true
from fastcore.basics import listify

from execnb.nbio import read_nb, write_nb

from nb_helpers.utils import (
    detect_imported_libs,
    git_main_name,
    search_string_in_nb,
    find_nbs,
    RichLogger,
    git_local_repo,
)
from nb_helpers.clean import clean_nb, clean_one
from nb_helpers.colab import add_colab_badge, add_colab_metadata, get_colab_url, has_colab_badge


WANDB_FEATURES = "Table,sweep,WandbCallback,WandbLogger,Artifact"
PYTHON_LIBS = "torch,keras,tensorflow,sklearn,yolo,jax,pandas,numpy,spacy,transformers,lightning,fastai"


def get_wandb_tracker(nb):
    "Get the value inside <!--- @wandbcode{tracker} -->"
    for cell in nb["cells"]:
        if "@wandbcode" in cell["source"]:
            tracker_id = re.search(r"@wandbcode{(.*?)}", cell["source"]).group(1)
            return tracker_id.split(",")[0]
    return ""


def search_code(nb, features=WANDB_FEATURES):
    "Search notebook for features"
    present_features = []
    for feat in listify(features.split(",")):
        if search_string_in_nb(nb, feat):
            present_features.append(feat)
    return present_features


@call_parse
def summary_nbs(
    path: Param("A path to nb files", Path, nargs="?", opt=False) = os.getcwd(),
    wandb_features: Param("wandb features to identify, comma separated", str) = WANDB_FEATURES,
    out_file: Param("Export to csv file", Path) = "summary_table.csv",
    github_issue: Param("Creates a `github_issue.md` file ready to be put online", store_true) = True,
    full_path: Param("Use full path for fname", store_true) = False,
):
    path = Path(path)
    # out_file = (path.parent / out_file).with_suffix(".csv")
    logger = RichLogger(columns=["fname", "tracker", "wandb", "python libs", "colab_idx"], out_file=out_file)

    files = find_nbs(path)
    assert len(files) > 0, "There is no `ipynb` notebooks in the path you submited"

    logger.info(f"Reading {len(files)} notebooks")

    repo_path = git_local_repo(files[0])

    for nb_path in files:
        nb = read_nb(nb_path)
        tracker_id = get_wandb_tracker(nb)
        fname = nb_path.name if not full_path else nb_path.relative_to(repo_path)
        features = search_code(nb, wandb_features)
        libs = detect_imported_libs(nb)
        colab_cell_idx = has_colab_badge(nb)
        row = [
            str(fname),
            tracker_id,
            ", ".join(features),
            ", ".join(libs),
            str(colab_cell_idx) if colab_cell_idx != -1 else "",
        ]
        colab_link = get_colab_url(nb_path, branch=git_main_name(nb_path))
        logger.writerow(row, colab_link)
    logger.to_table()
    logger.to_csv(Path(out_file).with_suffix(".csv"))
    logger.to_md(Path(out_file).with_suffix(".md"))

    if github_issue:
        logger.create_github_issue()


@call_parse
def fix_nbs(
    path: Param("A path to nb files", Path, nargs="?", opt=False) = os.getcwd(),
    colab_cell_idx: Param("Cell idx where the colab badge must go", int) = 0,
    branch: Param("The branch", str) = None,
):

    path = Path(path)
    files = find_nbs(path)
    assert len(files) > 0, "There is no `ipynb` notebooks in the path you submited"

    for nb_path in files:
        print(f"Add colab badge to {nb_path}")
        clean_one(nb_path, clear_outs=True)
        nb = read_nb(nb_path)
        tracker = get_wandb_tracker(nb)
        if tracker is not None and tracker != "":
            tracker = f"<!--- @wandbcode{{{tracker}}} -->"
        nb = add_colab_badge(nb, nb_path, branch=branch, idx=colab_cell_idx, tracker=tracker)
        add_colab_metadata(nb)
        clean_nb(nb)
        write_nb(nb, nb_path)


@call_parse
def filter_nbs(
    path: Param("A path to nb files", Path, nargs="?", opt=False) = os.getcwd(),
    lib_name: Param("Python lib names to filter, eg: tensorflow", str) = "",
    out_file: Param("Export to csv file", Path) = "filtered_nbs.csv",
):

    path = Path(path)

    logger = RichLogger(columns=["fname", "python libs"], out_file=out_file)

    files = find_nbs(path)
    assert len(files) > 0, "There is no `ipynb` notebooks in the path you submited"

    repo_path = git_local_repo(files[0])

    for nb_path in files:
        nb = read_nb(nb_path)
        fname = nb_path.relative_to(repo_path)
        libs = ", ".join(detect_imported_libs(nb))
        if lib_name in libs:
            row = [str(fname), libs]
            logger.writerow(row)

    logger.to_table()
    logger.to_csv(Path(out_file).with_suffix(".csv"), format_link=False)
