#!/usr/bin/env python
# -*- coding: utf-8 -*-

# Copyright 2025 EMBL - European Bioinformatics Institute
#
# 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.


import argparse
import os.path


def metadata_parser(drep_clstrs, derep_genomes):
    clusters = {}
    with open(drep_clstrs) as input_file:
        next(input_file)
        for line in input_file:
            l_line = line.rstrip().split(",")
            clstr_member = l_line[0].replace(".fa", "")
            clstr_num = "clstr_" + l_line[5]
            if clstr_num in clusters:
                clusters[clstr_num].append(clstr_member)
            else:
                clusters[clstr_num] = [clstr_member]

    reps_clusters = {}
    rep_assemblies_list = os.listdir(derep_genomes)
    rep_genomes_list = [f.replace(".fa", "") for f in rep_assemblies_list]
    for clstr in clusters:
        members = clusters[clstr]
        for genome in members:
            if genome in rep_genomes_list:
                reps_clusters[genome] = members

    return reps_clusters


def gff_gather(reps_clusters, gff_files):
    gff_files_abs = os.path.abspath(gff_files)
    cluster_counter = 0
    total_clusters = len(reps_clusters)

    for rep in reps_clusters:
        cluster_counter += 1
        print(
            "Processing cluster number "
            + str(cluster_counter)
            + " size "
            + str(len(reps_clusters[rep]))
            + " out of "
            + str(total_clusters)
        )

        if len(reps_clusters[rep]) > 1:
            dest_path = "rep_" + rep
            os.makedirs(dest_path, exist_ok=True)

            for member in reps_clusters[rep]:
                src_file = os.path.join(gff_files_abs, member + ".gff")
                link_name = os.path.join(dest_path, member + ".gff")

                if os.path.exists(src_file):
                    os.symlink(src_file, link_name)
                    print(f"Created symlink: {link_name} -> {src_file}")
                else:
                    print(f"Source file does not exist: {src_file}")


def main():
    parser = argparse.ArgumentParser(
        description="This script prepare the inputs to launch panaroo on custom databases"
    )
    parser.add_argument(
        "--drep_clstrs",
        type=str,
        help="Data table generated by drep: drep_output/data_tables/Cdb.csv",
        required=True,
    )
    parser.add_argument(
        "--gff_files",
        type=str,
        help="Path to the gff files to be used as input for panaroo. Files should be named as $genome.gff",
        required=True,
    )
    parser.add_argument(
        "--derep_genomes",
        type=str,
        help="Path to the dereplicated genomes: drep_output/dereplicated_genomes/",
        required=True,
    )
    args = parser.parse_args()

    ### Calling functions
    reps_clusters = metadata_parser(args.drep_clstrs, args.derep_genomes)
    gff_gather(reps_clusters, args.gff_files)


if __name__ == "__main__":
    main()
