# Copyright 2018 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Filesystem DB: poor man's worker coordination strategy.

This module works equivalently with local filesystems and GCS.
"""
import os

from absl import flags
from tensorflow import gfile
import re

import shipname

flags.DEFINE_string(
    'base_dir', None,
    'Root directory if using local FS as the database.'
    'Leave blank if using bucket_name.')

flags.DEFINE_string(
    'bucket_name', None,
    'Bucket name if using GCS as the filesystem DB.'
    'Leave blank is using base_dir.')

flags.register_multi_flags_validator(
    ['base_dir', 'bucket_name'],
    lambda flags: bool(flags['base_dir']) != bool(flags['bucket_name']),
    'Exactly one of --base_dir, --bucket_name must be set!')

FLAGS = flags.FLAGS


def _with_base(*args):
    def inner():
        base_dir = FLAGS.base_dir or 'gs://{}'.format(FLAGS.bucket_name)
        return os.path.join(base_dir, *args)
    return inner


# Functions to compute various important directories, based on FLAGS input.
models_dir = _with_base('models')
selfplay_dir = _with_base('data', 'selfplay')
holdout_dir = _with_base('data', 'holdout')
sgf_dir = _with_base('sgf')
eval_dir = _with_base('sgf', 'eval')
training_chunk_dir = _with_base('data', 'training_chunks')
golden_chunk_dir = _with_base('data', 'golden_chunks')


def get_models():
    """Finds all models, returning a list of model number and names
    sorted increasing.

    Returns: [(13, 000013-modelname), (17, 000017-modelname), ...etc]
    """
    all_models = gfile.Glob(os.path.join(models_dir(), '*.meta'))
    model_filenames = [os.path.basename(m) for m in all_models]
    model_numbers_names = sorted([
        (shipname.detect_model_num(m), shipname.detect_model_name(m))
        for m in model_filenames])
    return model_numbers_names


def get_latest_model():
    """Finds the latest model, returning its model number and name

    Returns: (17, 000017-modelname)
    """
    return get_models()[-1]


def get_model(model_num):
    """Given a model number 17, returns its full name 000017-modelname."""
    model_names_by_num = dict(get_models())
    return model_names_by_num[model_num]


def get_hour_dirs(root=None):
    """Gets the directories under selfplay_dir that match YYYY-MM-DD-HH."""
    root = root or selfplay_dir()
    return list(filter(lambda s: re.match(r"\d{4}-\d{2}-\d{2}-\d{2}", s),
                       gfile.ListDirectory(root)))


def game_counts(n_back=20):
    """Prints statistics for the most recent n_back models"""
    for _, model_name in get_models[-n_back:]:
        games = gfile.Glob(os.path.join(selfplay_dir(), model_name, '*.zz'))
        print("Model: {}, Games: {}".format(model_name, len(games)))