import gffutils
import sys
import os
import subprocess
from collections import defaultdict
import re
from lifton import extract_sequence


class Annotation():

    def __init__(self, file_name, infer_genes, infer_transcripts, merge_strategy, id_spec, force, verbose, auto_convert_gtf=True):
        self.file_name = file_name
        self.merge_strategy = merge_strategy
        self.id_spec = id_spec
        self.force = force
        self.verbose = verbose
        self.auto_convert_gtf = auto_convert_gtf
        
        # Detect file format and adjust inference settings
        # Default is GFF3 format (most common and recommended)
        file_format = self.detect_file_format()
        
        if file_format == "GTF format":
            print(f"INFO: Detected GTF format for {file_name}")
            print("INFO: GTF format support is experimental. For best results, consider converting to GFF3 using:")
            print("  agat_sp_gtf2gff.pl --gtf input.gtf -o output.gff3")
            print("  or")
            print("  gffread -E input.gtf -o output.gff3")
            # For GTF files, we need to infer genes and transcripts
            # GTF files don't have explicit gene/transcript features, only exons/CDS with gene_id/transcript_id
            # So we always enable inference for GTF files (user can't really override this)
            self.infer_genes = True
            self.infer_transcripts = True
            if self.verbose:
                print("  Automatically enabling gene and transcript inference for GTF format")
            
            # Optionally convert GTF to GFF3 for better compatibility
            if self.auto_convert_gtf:
                converted_file = self.convert_gtf_to_gff3()
                if converted_file:
                    # Verify the converted file exists and is valid
                    if os.path.exists(converted_file) and os.path.getsize(converted_file) > 0:
                        self.file_name = converted_file
                        # Re-detect format to ensure it's now GFF3
                        converted_format = self.detect_file_format()
                        if converted_format == "GFF format":
                            if self.verbose:
                                print(f"  Successfully converted GTF to GFF3: {converted_file}")
                        else:
                            if self.verbose:
                                print(f"  Warning: Converted file may not be valid GFF3 format")
                    else:
                        if self.verbose:
                            print(f"  Warning: Conversion produced empty or invalid file, using original GTF")
        elif file_format == "Unknown format":
            # If format detection fails, default to GFF3 and use user-provided settings
            if self.verbose:
                print(f"Warning: Could not determine file format for {file_name}, assuming GFF3 format")
            self.infer_genes = infer_genes
            self.infer_transcripts = infer_transcripts
        else:
            # For GFF3 files (or any other recognized format), use user-provided settings
            # argparse passes False when flags are not provided, True when they are
            self.infer_genes = infer_genes
            self.infer_transcripts = infer_transcripts
        
        self.get_db_connnection()
        
    def get_db_connnection(self):
        try:
            feature_db = gffutils.FeatureDB(self.file_name)
        except:
            feature_db = self.build_database()
        self.db_connection = feature_db


    def build_database(self):
        if self.infer_genes:
            disable_genes = False
        else:
            disable_genes = True
        if self.infer_transcripts:
            disable_transcripts = False
        else:
            disable_transcripts = True
        try:
            transform_func = self.get_transform_func()
            feature_db = gffutils.create_db(self.file_name, self.file_name + "_db", 
                                        merge_strategy=self.merge_strategy, 
                                        id_spec=self.id_spec,
                                        force=self.force,
                                        verbose=self.verbose, disable_infer_transcripts=disable_transcripts,
                                            disable_infer_genes=disable_genes, transform=transform_func)
        except Exception as e:
            error_msg = str(e)
            if "UNIQUE constraint failed" in error_msg or "duplicate" in error_msg.lower():
                print(f"Warning: Duplicate feature IDs detected in {self.file_name}. Attempting to handle duplicates...")
            else:
                print(f"gffutils database build failed with: {e}")
            feature_db = self.build_database_again()
            return feature_db


    def build_database_again(self):
        if self.infer_genes:
            disable_genes = False
        else:
            disable_genes = True
        if self.infer_transcripts:
            disable_transcripts = False
        else:
            disable_transcripts = True
        # For duplicate IDs (especially non-overlapping CDS features), use unique ID transformation
        # This is better than 'merge' strategy which is designed for overlapping features
        try:
            transform_func = self.get_unique_id_transform()
            feature_db = gffutils.create_db(self.file_name, self.file_name + "_db", 
                                        merge_strategy="create_unique",
                                        id_spec=self.id_spec,
                                        force=True,
                                        verbose=self.verbose, disable_infer_transcripts=disable_transcripts,
                                            disable_infer_genes=disable_genes, transform=transform_func)
            print("Successfully built database using unique ID transformation to handle duplicate IDs.")
            return feature_db
        except Exception as e1:
            print(f"Unique ID transformation failed: {e1}")
            print("Attempting to build database with merge strategy...")
            # If transform fails, try merge strategy as a fallback
            try:
                transform_func = self.get_transform_func()
                feature_db = gffutils.create_db(self.file_name, self.file_name + "_db", 
                                            merge_strategy="merge",
                                            id_spec=self.id_spec,
                                            force=True,
                                            verbose=self.verbose, disable_infer_transcripts=disable_transcripts,
                                                disable_infer_genes=disable_genes, transform=transform_func)
                print("Successfully built database using merge strategy to handle duplicate IDs.")
                return feature_db
            except Exception as e2:
                print(f"ERROR: Failed to build database after trying multiple strategies.")
                print(f"Last error: {e2}")
                print(f"Please check the GFF file {self.file_name} for duplicate feature IDs or other issues.")
                sys.exit(1)


    def get_transform_func(self):
        if self.infer_genes is False:
            return None
        else:
            return transform_func

    def get_unique_id_transform(self):
        """Create a transform function that ensures unique IDs by appending a counter to duplicates.
        This is especially important for non-overlapping CDS features in RefSeq annotations that may have duplicate IDs.
        """
        # Dictionary to track how many times we've seen each ID
        seen_ids = {}
        
        def unique_id_transform(feature):
            # Apply the original transform if needed
            if self.infer_genes:
                feature = transform_func(feature)
            
            # Get the ID from the feature - gffutils uses the 'id' property or ID attribute
            # We need to check both the id property and the ID attribute
            original_id = None
            if hasattr(feature, 'id') and feature.id:
                original_id = feature.id
            elif 'ID' in feature.attributes and len(feature.attributes['ID']) > 0:
                original_id = feature.attributes['ID'][0]
            
            # If no ID, gffutils will generate one, so we don't need to modify
            if not original_id:
                return feature
            
            # Track seen IDs and make unique if needed
            # First occurrence keeps original ID, subsequent ones get _dup1, _dup2, etc.
            if original_id in seen_ids:
                seen_ids[original_id] += 1
                new_id = f"{original_id}_dup{seen_ids[original_id]}"
                # Update both the ID attribute and the id property
                # This ensures gffutils uses the new unique ID
                if 'ID' in feature.attributes:
                    feature.attributes['ID'] = [new_id]
                if hasattr(feature, 'id'):
                    feature.id = new_id
            else:
                # First time seeing this ID - mark it as seen but don't modify
                seen_ids[original_id] = 0
            
            return feature
        
        return unique_id_transform



    def get_protein_coding_features(self, feature_types):
        protein_coding_genes = []
        for f_type in feature_types:
            for feature in self.db_connection.features_of_type(featuretype=f_type):
                if self.is_highest_parent(feature):
                    CDS_list = [child for child in self.db_connection.children(feature) if child.featuretype == 'CDS']
                    if len(CDS_list) > 0:
                        protein_coding_genes.append(feature.id)
        return protein_coding_genes


    def get_noncoding_features(self, feature_types):
        noncoding_genes = []
        for f_type in feature_types:
            for feature in self.db_connection.features_of_type(featuretype=f_type):
                if self.is_highest_parent(feature):
                    CDS_list = [child for child in self.db_connection.children(feature) if child.featuretype == 'CDS']
                    if len(CDS_list)==  0:
                        noncoding_genes.append(feature.id)
        return noncoding_genes


    def get_novel_protein_coding_features(self, ref_genes, feature_types):
        novel_protein_coding = []
        all_protein_coding = self.get_protein_coding_features(feature_types)
        for feature in all_protein_coding:
            if feature not in ref_genes:
                novel_protein_coding.append(feature)
        return novel_protein_coding


    def get_novel_noncoding_features(self,ref_genes, feature_types):
        novel_noncoding = []
        all_noncoding = self.get_noncoding_features(feature_types)
        for feature in all_noncoding:
            if feature not in ref_genes:
                novel_noncoding.append(feature)
        return novel_noncoding


    def get_all_parent_feature_ids(self, feature_types):
        feature_ids = []
        for f_type in feature_types:
            feature_ids += [feature.id for feature in self.db_connection.features_of_type(
                featuretype=f_type) if self.is_highest_parent(feature)]
        return feature_ids


    def make_parent_to_child_dict(self, protein_coding, gene):
        child_dict=defaultdict(list)
        child_list = self.db_connection.children(gene)
        child_count = 0
        for child in child_list:
            child_count += 1
            if (protein_coding and child.featuretype == "CDS") or (protein_coding is False and self.is_lowest_child(
                    child)):
                parent = list(self.db_connection.parents(child, level=1))[0]
                child_dict[parent.id].append(child)
        if child_count == 0:
            child_dict[gene] = [self.db_connection[gene]]
        return child_dict


    def get_features_of_type(self,feature_types):
        features_of_type = []
        for f_type in feature_types:
            features_of_type += list(self.db_connection.features_of_type(f_type))
        return features_of_type


    def get_feature_dict(self, feature_types):
        id_to_feature = {}
        features = self.get_features_of_type(feature_types)
        for feature in features:
            id_to_feature[feature.id] = feature
        return id_to_feature


    def get_source_name(self, feature_name):
        feature = self.db_connection[feature_name]
        if 'extra_copy_number' in feature.attributes:
            if feature.attributes['extra_copy_number'][0] == feature.id.split("_")[-1]:
                return "_".join(feature.id.split("_")[:-1])
        return feature.id


    def is_lowest_child(self, feature_name):
        return len(list(self.db_connection.children(feature_name))) == 0


    def is_highest_parent(self, feature_name):
        return len(list(self.db_connection.parents(feature_name))) == 0


    def get_paralog_name(self, feature_name):
        gene_attributes = self.db_connection[feature_name].attributes
        if "extra_copy_number" in gene_attributes:
            paralog_name = "_".join(gene_attributes["ID"][0].split("_")[:-1])
            return paralog_name
        return ""


    def get_num_levels(self, feature_name):
        level = 1
        new_children = list(self.db_connection.children(feature_name, level=level))
        while len(new_children) !=0:
            level += 1
            new_children = list(self.db_connection.children(feature_name, level=level))
        return level

    def detect_file_format(self):
        """Detect if the input file is GTF or GFF3 format."""
        try:
            return extract_sequence.determine_file_format(self.file_name)
        except Exception as e:
            if self.verbose:
                print(f"Warning: Could not detect file format: {e}")
            # Default to GFF3 if detection fails
            return "GFF format"

    def convert_gtf_to_gff3(self):
        """Convert GTF file to GFF3 format using agat or gffread.
        Returns the path to the converted file, or None if conversion fails.
        """
        # Check if agat is available (try different agat commands)
        agat_available = False
        agat_cmd = None
        for cmd_name in ["agat_convert_sp_gff2gtf.pl", "agat_sp_gtf2gff.pl", "agat"]:
            if self._check_tool_available(cmd_name):
                agat_available = True
                agat_cmd = cmd_name
                break
        
        gffread_available = self._check_tool_available("gffread")
        
        if not agat_available and not gffread_available:
            if self.verbose:
                print("  Warning: Neither agat nor gffread found. GTF file will be used directly.")
                print("  For best results with GTF files, consider converting to GFF3 using:")
                print("    agat_sp_gtf2gff.pl --gtf input.gtf -o output.gff3")
                print("  or")
                print("    gffread -E input.gtf -o output.gff3")
            return None
        
        # Create output filename
        base_name = os.path.splitext(self.file_name)[0]
        output_file = base_name + "_converted.gff3"
        
        # Try gffread first (more commonly available and reliable)
        if gffread_available:
            try:
                cmd = ["gffread", "-E", self.file_name, "-o", output_file]
                result = subprocess.run(cmd, capture_output=True, text=True, check=True)
                if os.path.exists(output_file) and os.path.getsize(output_file) > 0:
                    return output_file
            except (subprocess.CalledProcessError, FileNotFoundError) as e:
                if self.verbose:
                    print(f"  Warning: gffread conversion failed: {e}")
        
        # Try agat as fallback
        if agat_available:
            try:
                # Try different agat command formats
                if "gtf2gff" in agat_cmd:
                    cmd = [agat_cmd, "--gtf", self.file_name, "-o", output_file]
                else:
                    # Generic agat command - may need adjustment
                    cmd = [agat_cmd, "--gff", self.file_name, "-o", output_file]
                result = subprocess.run(cmd, capture_output=True, text=True, check=True)
                if os.path.exists(output_file) and os.path.getsize(output_file) > 0:
                    return output_file
            except (subprocess.CalledProcessError, FileNotFoundError) as e:
                if self.verbose:
                    print(f"  Warning: agat conversion failed: {e}")
        
        return None

    def _check_tool_available(self, tool_name):
        """Check if a command-line tool is available."""
        try:
            subprocess.run(["which", tool_name], capture_output=True, check=True)
            return True
        except (subprocess.CalledProcessError, FileNotFoundError):
            return False



def transform_func(x):
    if 'transcript_id' in x.attributes:
        x.attributes['transcript_id'][0] += '_transcript'
    return x


def get_perc_id(feature):
    if 'sequence_ID' in feature.attributes:
        return float(feature.attributes['sequence_ID'][0])
    return 0.0


def merge_children_intervals(children):
    if len(children) == 0:
        return []
    intervals = [[child.start, child.end] for child in children]
    intervals.sort(key=lambda interval: interval[0])
    merged = [intervals[0]]
    for current in intervals:
        previous = merged[-1]
        if current[0] <= previous[1]:
            previous[1] = max(previous[1], current[1])
        else:
            merged.append(current)
    return merged