from typing import List, Optional

import os
import dotenv

from higgsfield.internal.cfg import AppConfig

from .params import Param, build_gh_action_inputs, build_run_params
from jinja2 import Environment, FileSystemLoader, Template
from higgsfield.internal.util import templates_path
from pathlib import Path
from importlib.machinery import SourceFileLoader
from .decorator import ExperimentDecorator
from higgsfield.internal.experiment.ast_parser import parse_experiments

header = """# THIS FILE WAS GENERATED BY HIGGSFIELD. 
# DO NOT EDIT. 
# IF YOUR WORKFLOW DOESN'T WORK, CREATE AN ISSUE.
"""


class ActionBuilder:
    app_config: AppConfig
    wf_dir: Path
    template: Template
    template_name: str

    def __init__(self, app_config: AppConfig, project_path: Path):
        environment = Environment(loader=FileSystemLoader(templates_path()))
        self.template = environment.get_template(self.template_name)

        self.app_config = app_config

        wf_dir = project_path / ".github" / "workflows"
        if not wf_dir.exists():
            wf_dir.mkdir(parents=True, exist_ok=True)
        elif not wf_dir.is_dir():
            raise ValueError(f"{wf_dir} is not a directory, delete it and try again")

        self.wf_dir = wf_dir


class KillBuilder(ActionBuilder):
    template_name: str = "kill_action.j2"

    def generate(self):
        (self.wf_dir / "kill.yml").write_text(
            self.template.render(
                header=header,
                project_name=self.app_config.name,
            )
        )

        print("Updated kill action")


def as_keyed_repo_url(repo_url: Optional[str], project_name: str) -> str:
    if repo_url is None:
        raise ValueError("Did you add git remote origin? / push to github?")

    # get the index of "/" after "github.com"
    # and replace it with "-project_name:"
    return repo_url.replace("github.com:", f"github.com-{project_name}:")


def insert_env_line(keys: List[str], indent: str) -> str:
    lines = []
    for key in keys:
        line = indent + "echo " + key + '="${{ secrets.' + key + ' }}" >> env'
        lines.append(line)

    return "\n".join(lines)


def env_keys_as_action(path: Path, indent: str) -> str:
    keys = list(dotenv.dotenv_values(path).keys())
    keys.pop(keys.index("SSH_KEY"))

    return insert_env_line(keys, indent)


echo_indent = "          "


class DeployBuilder(ActionBuilder):
    template_name = "deploy_action.j2"

    def generate(self):
        (self.wf_dir / "deploy.yml").write_text(
            self.template.render(
                header=header,
                project_name=self.app_config.name,
                keyed_repo_url=as_keyed_repo_url(
                    self.app_config.github_repo_url, self.app_config.name
                ),
                env_gen=env_keys_as_action(
                    self.wf_dir.parent.parent / "env", echo_indent
                ),
            )
        )

        print("Updated deploy action")


class ExperimentBuilder(ActionBuilder):
    template_name = "experiment_action.j2"

    def generate(self, experiment_name: str, params: List[Param]):
        (self.wf_dir / f"run_{experiment_name}.yml").write_text(
            self.template.render(
                header=header,
                experiment_name=experiment_name,
                project_name=self.app_config.name,
                params=build_gh_action_inputs(params),
                rest=build_run_params(params),
                env_gen=env_keys_as_action(
                    self.wf_dir.parent.parent / "env", echo_indent
                ),
            )
        )
        

        print("Updated experiment action", experiment_name)


def _source_experiments(base_path: Path):
    """
    Only used inside docker to inject experiments into the module.
    Do not use outside of docker. But if you do, you will have to have
    the same dependencies (aka environment) as the docker container.
    """
    for file in base_path.glob("**/*.py"):
        module_name = os.path.basename(file).split(".py")[0].split(".py")[0].replace(" ", "_").replace("-", "_")
        SourceFileLoader(module_name, str(file)).load_module()


def build_all_experiment_actions(wd_path: Path, app_config: AppConfig):
    """
    Builds all experiment actions

    Root path should be the root of the project and must contain the src folder under which everything is defined.
    Project name should be the name of the project.

    It will parse ast of all files under src folder and search for experiments.
    Then it will build actions for each experiment and save them under .github/workflow folder.
    Deletes all actions that have name prefix run_experiment_ and header inside if not needed.
    """
    exp_params_pairs = []
    for file in (wd_path / "src").glob("**/*.py"):
        exp_params_pairs.extend(parse_experiments(str(file.resolve())))

    if len(exp_params_pairs) == 0:
        print("No experiments found")
        return

    experiments = ExperimentDecorator.from_ast(exp_params_pairs)
    actions_folder = wd_path / ".github" / "workflows"
    actions_folder.mkdir(parents=True, exist_ok=True)

    # list all files that have name prefix run_experiment_ and header inside
    for i in actions_folder.glob("run_*.yml"):
        if i.read_text().startswith(header):
            i.unlink()

    for experiment_name, experiment in experiments.items():
        ExperimentBuilder(app_config, wd_path).generate(
            experiment_name, experiment.params
        )
