#!/usr/bin/env python3
"""
Tool: fastq_qc_trimmer.py
Category: High-Throughput Sequencing & Quality Control
Description: High-performance streaming FASTQ trimmer with sliding-window Phred quality filtering,
             3' adapter removal, and minimum length retention. Handles both plain and .gz files.
Usage:
    python fastq_qc_trimmer.py -i reads.fastq.gz -o trimmed.fastq.gz --min-q 25 --window 4 --min-len 35
    python fastq_qc_trimmer.py -i input.fastq --adapter AGATCGGAAGAGC
"""

import argparse
import gzip
import sys

def open_fastq(path, mode='rt'):
    if path.endswith('.gz'):
        return gzip.open(path, mode)
    return open(path, mode)

def trim_sliding_window(seq: str, qual: str, window_size: int = 4, q_threshold: int = 20, offset: int = 33) -> tuple:
    """Performs 3' sliding window trimming based on average Phred score."""
    scores = [ord(c) - offset for c in qual]
    trim_pos = len(scores)
    
    for i in range(len(scores) - window_size, -1, -1):
        window_scores = scores[i:i + window_size]
        avg_q = sum(window_scores) / window_size
        if avg_q < q_threshold:
            trim_pos = i
        else:
            break
            
    return seq[:trim_pos], qual[:trim_pos]

def clip_adapter(seq: str, qual: str, adapter: str) -> tuple:
    """Clips 3' sequencing adapter if detected with exact prefix match."""
    idx = seq.find(adapter)
    if idx != -1:
        return seq[:idx], qual[:idx]
    return seq, qual

def run_trimmer(in_path, out_path, min_q=20, window_size=4, min_len=30, adapter=None):
    total = 0
    passed = 0
    trimmed_bases = 0
    
    out_fh = open_fastq(out_path, 'wt') if out_path else None
    
    with open_fastq(in_path, 'rt') as f:
        while True:
            h = f.readline().strip()
            if not h:
                break
            seq = f.readline().strip()
            c = f.readline().strip()
            qual = f.readline().strip()
            total += 1
            orig_len = len(seq)
            
            # Clip adapter
            if adapter:
                seq, qual = clip_adapter(seq, qual, adapter)
                
            # Sliding window trim
            seq, qual = trim_sliding_window(seq, qual, window_size=window_size, q_threshold=min_q)
            
            trimmed_bases += (orig_len - len(seq))
            
            if len(seq) >= min_len:
                passed += 1
                if out_fh:
                    out_fh.write(f"{h}\n{seq}\n{c}\n{qual}\n")
                    
    if out_fh:
        out_fh.close()
        
    return total, passed, trimmed_bases

def main():
    parser = argparse.ArgumentParser(description="Streaming FASTQ quality trimmer and adapter filter.")
    parser.add_argument("-i", "--input", required=False, help="Input FASTQ (.fastq or .fastq.gz)")
    parser.add_argument("-o", "--output", help="Output trimmed FASTQ path")
    parser.add_argument("-q", "--min-q", type=int, default=20, help="Phred quality cutoff for sliding window (default: 20)")
    parser.add_argument("-w", "--window", type=int, default=4, help="Sliding window size (default: 4)")
    parser.add_argument("-l", "--min-len", type=int, default=30, help="Minimum read length to retain (default: 30)")
    parser.add_argument("-a", "--adapter", default=None, help="3' adapter sequence to clip (e.g. AGATCGGAAGAGC)")
    args = parser.parse_args()

    if not args.input:
        # Create demo fastq
        demo = "/tmp/demo_trim_in.fastq"
        with open(demo, "w") as f:
            f.write("@READ_1\nATCGATCGATCGATCGATCGAGATCGGAAGAGC\n+\nIIIIIIIIIIIIIIIIIIII#############\n")
            f.write("@READ_2\nAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA\n+\nIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIII\n")
        args.input = demo
        print(f"[!] Using demo FASTQ at {demo}\n")

    total, passed, tr_bases = run_trimmer(args.input, args.output, args.min_q, args.window, args.min_len, args.adapter)
    print(f"=== FASTQ Trimming Results ===")
    print(f"Total Input Reads : {total:,}")
    print(f"Retained Reads    : {passed:,} ({(passed/total*100):.2f}%)")
    print(f"Discarded Reads   : {total - passed:,}")
    print(f"Total Bases Trimmed: {tr_bases:,}")

if __name__ == "__main__":
    main()