from typing import Iterator
from collections import defaultdict
from google.cloud import storage

from ruamel.yaml import YAML


parse_sources = [
    {
        'name': 'metadb_new',
        'source_project': 'gcp-digital-bi',
        'target_project': 'gcp-digital-bi-dev',
        'gcs_bucket': "gs://metadb",
        'gcs_prefix': "metadb/"
    },
]

def parse_gcs_files(project: str, bucket: str, prefix: str) -> Iterator[str]:
    client = storage.Client(project=project)
    bucket = client.get_bucket(bucket.replace("gs://", ""))
    blobs = bucket.list_blobs(prefix=prefix)
    for blob in blobs:
        yield blob.name

output_dict = {'sources': []}

def extract_directories(files: Iterator[str]) -> dict:
    directories = defaultdict(lambda: {})
    for file in files:
        if file.startswith('metadb/cost'):
            continue
        parts = file.split('/')
        if len(parts) > 2:
            path = '/'.join(parts[1:-2])
            name = parts[-2]
            if path in ('dedicated', 'serverless'):
                name = f"{name}_{path}"
            
            gcs_glob = f"{'/'.join(parts[0:-1])}/*"

            if name in directories:
                if directories[name]["path"] != path or directories[name]["gcs_glob"] != gcs_glob:
                    raise ValueError(f"Inconsistent attributes for {name}, in {file}\n {path}, {directories[name]["path"]}; {gcs_glob}, {directories[name]["gcs_glob"]}")
            else:
                directories[name]["name"] = name
                directories[name]["identifier"] = name
                directories[name]["gcs_glob"] = gcs_glob
                directories[name]["path"] = path
    
    return list(directories.values())


yaml = YAML()

for parse_source in parse_sources:
    files = parse_gcs_files(
        parse_source['source_project'],
        parse_source['gcs_bucket'],
        parse_source['gcs_prefix']
    )
    source = {
        'name': parse_source['name'],
        'schema': f"raw_{parse_source['name']}",
        'database': parse_source['target_project'],
        'source_bucket': parse_source['gcs_bucket'],
        'source_project': parse_source['source_project'],
        'tables': extract_directories(files)
    }
    output_dict['sources'].append(source)



with open('output.yml', 'w') as fout:
    yaml.dump(output_dict, fout)