import logging
import json
import os
import requests
import sys
import yaml
from typing import Optional

EXPECTED_TF_CONFIG_FILE_NAMES = ["tdfconfig.yml", "tdfconfig.yaml", "validate_configs.yaml", "commit_configs.yaml"]
TRANSFORM_API_URL = "web-api.prod.transformdata.io"
UPLOAD_MODE_VALIDATE = "validate"
UPLOAD_MODE_COMMIT = "commit"


logger = logging.getLogger(__name__)


def read_config_files(config_dir: str):
    """Read yaml files from config_dir. Returns (file name, file contents) per file in dir"""
    assert os.path.exists(config_dir), f"User-specified config dir ({config_dir}) does not exist"

    results = {}
    for path, _folders, filenames in os.walk(config_dir):
        for fname in filenames:
            if not (fname.endswith(".yml") or fname.endswith(".yaml")):
                continue

            # ignore transform config
            if fname in EXPECTED_TF_CONFIG_FILE_NAMES:
                continue

            with open(os.path.join(path, fname), "r") as f:
                filepath = os.path.abspath(os.path.join(path, fname))
                results[filepath] = f.read()
                try:
                    yaml.safe_load_all(results[filepath])
                except yaml.parser.ParserError as e:
                    raise yaml.parser.ParserError(f"Invalid yaml in config file at path: {filepath}. Exception: '{e}'")
                except Exception as e:
                    raise Exception(f"Failed loading yaml config file. Error {e}")

                results[fname] = f.read()

    return results


def upload_configs(  # noqa: D
    mode: str,
    api_key: str,
    repo: Optional[str] = None,
    branch: Optional[str] = None,
    commit: Optional[str] = None,
    config_dir: Optional[str] = None,
    api_url: Optional[str] = None,
) -> requests.Response:
    yaml_files = read_config_files(config_dir or ".")
    results = {"yaml_files": yaml_files}
    logger.info(f"Files to upload: {yaml_files.keys()}")
    headers = {"Content-Type": "application/json", "Authorization": f"X-Api-Key {api_key}"}

    add_files_url = f"https://{api_url or TRANSFORM_API_URL}/api/v1/model/{repo}/{branch}/{commit}/add_model_files"
    logger.info(f"add_files_url: {add_files_url}")
    logger.info("Uploading config files")
    r = requests.post(add_files_url, data=json.dumps(results).encode("utf-8"), headers=headers)
    if r.status_code != 200:
        raise Exception(f"Failed uploading config yaml files. Error: {r.text}")

    if mode == UPLOAD_MODE_VALIDATE:
        validate_url = f"https://{api_url or TRANSFORM_API_URL}/api/v1/model/{repo}/{branch}/{commit}/validate_model"
        logger.info(f"validate_url: {validate_url}")
        logger.info("Checking that uploaded configs are valid and form a valid model")
        r = requests.post(validate_url, headers=headers)
        if r.status_code != 200:
            raise Exception(f"Failed validating uploaded configs. Error response: {r.text}")
        logger.info("Successfully validated configs")
    elif mode == UPLOAD_MODE_COMMIT:
        commit_url = f"https://{api_url or TRANSFORM_API_URL}/api/v1/model/{repo}/{branch}/{commit}/commit_model"
        logger.info(f"commit_url: {commit_url}")
        logger.info("Committing model")
        r = requests.post(commit_url, headers=headers)
        if r.status_code != 200:
            raise Exception(f"Failed committing uploaded configs. Error response: {r.text}")
        logger.info("Successfully committed configs")
    else:
        raise ValueError(f"Invalid upload mode ({mode}). Expected '{UPLOAD_MODE_VALIDATE}' or '{UPLOAD_MODE_COMMIT}'")

    return r


if __name__ == "__main__":
    mode = sys.argv[1]
    if mode != UPLOAD_MODE_VALIDATE and mode != UPLOAD_MODE_COMMIT:
        raise ValueError(f"Invalid upload mode ({mode}) passed via args.")

    # Retrieve git info and API key from env
    REPO = os.getenv("REPO")
    # remove github org from repo
    if REPO:
        REPO = "/".join(REPO.split("/")[1:])

    if os.getenv("GITHUB_HEAD_REF") == "":
        BRANCH = os.getenv("GITHUB_REF").lstrip("/refs/heads/")
    else:
        BRANCH = os.getenv("GITHUB_HEAD_REF")

    COMMIT = os.getenv("GITHUB_SHA")
    TRANSFORM_CONFIG_DIR = os.getenv("TRANSFORM_CONFIG_DIR").lstrip().rstrip()
    TRANSFORM_API_KEY = os.environ["TRANSFORM_API_KEY"].lstrip().rstrip()  # fail if TRANSFORM_API_KEY not present

    upload_configs(mode, TRANSFORM_API_KEY, REPO, BRANCH, COMMIT, TRANSFORM_CONFIG_DIR)
    print("success")
