#!/usr/bin/env python3
"""
Tool: gtf_gff_transcript_parser.py
Category: Genome Annotation & Feature Extraction
Description: Parses Ensembl/GENCODE GTF and GFF3 files, calculates exon/CDS spans,
             identifies the longest transcript isoform per gene, and exports BED6 tracks.
Usage:
    python gtf_gff_transcript_parser.py -i annotation.gtf -o transcripts.bed
"""

import argparse
import gzip
import sys

def parse_attributes(attr_str):
    """Parses semicolon-separated attributes in GTF format."""
    attrs = {}
    for part in attr_str.strip().split(';'):
        part = part.strip()
        if not part:
            continue
        if ' ' in part:
            k, v = part.split(' ', 1)
            attrs[k.strip()] = v.strip().strip('"')
    return attrs

def process_gtf(filepath, out_bed=None):
    opener = gzip.open if filepath.endswith('.gz') else open
    genes = {}
    
    with opener(filepath, 'rt') as f:
        for line in f:
            if line.startswith('#'):
                continue
            cols = line.strip().split('\t')
            if len(cols) < 9:
                continue
            chrom, source, feature, start, end, score, strand, frame, attr_str = cols
            
            if feature in ('exon', 'CDS'):
                attrs = parse_attributes(attr_str)
                gene_id = attrs.get('gene_id', 'UNKNOWN')
                trans_id = attrs.get('transcript_id', 'UNKNOWN')
                gene_name = attrs.get('gene_name', gene_id)
                
                span = int(end) - int(start) + 1
                if gene_id not in genes:
                    genes[gene_id] = {
                        'chrom': chrom, 'strand': strand, 'name': gene_name,
                        'transcripts': {}
                    }
                if trans_id not in genes[gene_id]['transcripts']:
                    genes[gene_id]['transcripts'][trans_id] = 0
                genes[gene_id]['transcripts'][trans_id] += span

    out_fh = open(out_bed, 'w') if out_bed else sys.stdout
    for gid, data in genes.items():
        # Find longest transcript
        best_tid, best_len = max(data['transcripts'].items(), key=lambda x: x[1])
        out_fh.write(f"{data['chrom']}\t{gid}\t{best_tid}\t{data['name']}\t{best_len} bp\t{data['strand']}\n")
        
    if out_bed and out_fh != sys.stdout:
        out_fh.close()
        
    return len(genes)

def main():
    parser = argparse.ArgumentParser(description="GTF/GFF3 Feature Extractor.")
    parser.add_argument("-i", "--input", help="Path to input GTF or GFF3")
    parser.add_argument("-o", "--output", help="Optional output path")
    args = parser.parse_args()

    if not args.input:
        demo = "/tmp/demo_annot.gtf"
        with open(demo, "w") as f:
            f.write('chr1\tHAVANA\texon\t1000\t1200\t.\t+\t.\tgene_id "ENSG001"; transcript_id "ENST001_A"; gene_name "ACTB";\n')
            f.write('chr1\tHAVANA\texon\t1500\t1800\t.\t+\t.\tgene_id "ENSG001"; transcript_id "ENST001_A"; gene_name "ACTB";\n')
            f.write('chr1\tHAVANA\texon\t1000\t1300\t.\t+\t.\tgene_id "ENSG001"; transcript_id "ENST001_B"; gene_name "ACTB";\n')
        args.input = demo
        print(f"[!] Using demo GTF at {demo}\n")

    n_genes = process_gtf(args.input, args.output)
    print(f"\n[Processed {n_genes} unique genes from GTF]", file=sys.stderr)

if __name__ == "__main__":
    main()