#' Strand-Wise Gene and Protein Extraction from GFF3 Files
#'
#' Extracts strand-wise genomic features, UTRs, CDS, introns,
#' exons, promoters, and associated protein sequences from a GFF3 file.
#'
#' @param GFF3 Path to GFF3 annotation file
#' @param PROMOTER_LEN Integer specifying promoter length
#' @param output_dir Directory where output CSV files will be saved
#'
#' @return Writes multiple CSV files to output directory
#'
#' @examples
#' gff_file <- system.file("extdata", "example.gff3", package = "GFFStrandLoc")
#' outdir <- tempdir()
#' \dontrun{
#' GFFStrandLoc(gff_file, 1000, outdir)
#' }
#'
#'@importFrom utils read.table tail write.csv
#' @export
GFFStrandLoc <- function(GFF3, PROMOTER_LEN, output_dir) {

  ## ==================== Setup ====================
  if (!dir.exists(output_dir)) {
    dir.create(output_dir, recursive = TRUE)
  }

  df <- read.table(
    GFF3,
    sep = "\t",
    comment.char = "#",
    header = FALSE,
    stringsAsFactors = FALSE,
    quote = ""
  )

  colnames(df) <- c(
    "ID","Source","Portion","Start","End",
    "score","Strand","Phase","Reference_id"
  )

  ## ==================== robust_gene_id ====================
  robust_gene_id <- function(ref, portion) {

    if (is.na(ref)) return(ref)

    if (grepl('gene_id "', ref)) {
      return(sub('.*gene_id "([^"]+)".*', '\\1', ref))
    }

    if (portion == "gene" && grepl("ID=", ref)) {
      return(sub(".*ID=([^;]+).*", "\\1", ref))
    }

    if (portion %in% c("mRNA","transcript") && grepl("Parent=", ref)) {
      return(sub(".*Parent=([^;]+).*", "\\1", ref))
    }

    if (grepl("Parent=", ref)) {
      parent <- sub(".*Parent=([^;]+).*", "\\1", ref)
      return(strsplit(parent, ",")[[1]][1])
    }

    if (grepl("\\.", ref)) {
      return(strsplit(ref, "\\.")[[1]][1])
    }

    return(ref)
  }

  df$Reference_id <- mapply(
    robust_gene_id,
    df$Reference_id,
    df$Portion
  )

  ## ==================== Protein sequences ====================
  protein_dict <- list()
  con <- file(GFF3, "r")
  current_gene <- NULL
  collecting <- FALSE
  seq <- character()

  while (length(line <- readLines(con, n = 1, warn = FALSE)) > 0) {

    line <- trimws(line)

    if (startsWith(line, "# start gene")) {
      current_gene <- tail(strsplit(line, " ")[[1]], 1)
      seq <- character()
      collecting <- FALSE

    } else if (startsWith(line, "# protein sequence")) {
      collecting <- TRUE
      seq <- c(seq, sub(".*\\[", "", line))

    } else if (collecting) {
      if (grepl("]", line)) {
        seq <- c(seq, gsub("]", "", line))
        protein_dict[[current_gene]] <- gsub(" ", "", paste(seq, collapse = ""))
        collecting <- FALSE
      } else {
        seq <- c(seq, line)
      }
    }
  }
  close(con)

  ## ---- SAFE protein mapper (NO LIST COLUMNS)
  safe_protein_map <- function(ids, prot_dict) {
    sapply(ids, function(x) {
      if (!is.null(prot_dict[[x]])) prot_dict[[x]] else NA_character_
    })
  }

  ## ==================== Strand split ====================
  Positive_strand <- df[df$Strand == "+", ]
  Negative_strand <- df[df$Strand == "-", ]

  add_len <- function(x) {
    x$length <- x$End - x$Start + 1
    x
  }

  ## ==================== POSITIVE STRAND ====================
  Positive_strand_transcript <- Positive_strand[Positive_strand$Portion=="transcript",]
  Positive_strand_start_codon <- Positive_strand[Positive_strand$Portion=="start_codon",]
  Positive_strand_stop_codon  <- Positive_strand[Positive_strand$Portion=="stop_codon",]

  merged <- merge(
    Positive_strand_transcript,
    Positive_strand_start_codon,
    by="Reference_id",
    suffixes=c("_tx","_sc")
  )

  five_prime_UTR_positive <- data.frame(
    ID=merged$ID_tx,
    Source=merged$Source_tx,
    Portion="five_prime_UTR",
    Start=merged$Start_tx,
    End=merged$Start_sc-1,
    score=merged$score_tx,
    Strand=merged$Strand_tx,
    Phase=".",
    Reference_id=merged$Reference_id,
    stringsAsFactors = FALSE
  )

  merged3 <- merge(
    Positive_strand_transcript,
    Positive_strand_stop_codon,
    by="Reference_id",
    suffixes=c("_tx","_sc")
  )

  Three_prime_UTR_positive <- data.frame(
    ID=merged3$ID_tx,
    Source=merged3$Source_tx,
    Portion="three_prime_UTR",
    Start=merged3$End_sc+1,
    End=merged3$End_tx,
    score=merged3$score_tx,
    Strand=merged3$Strand_tx,
    Phase=".",
    Reference_id=merged3$Reference_id,
    stringsAsFactors = FALSE
  )

  CDS_Positive_strand    <- add_len(Positive_strand[Positive_strand$Portion=="CDS",])
  intron_positive_strand <- add_len(Positive_strand[Positive_strand$Portion=="intron",])
  exon_positive_strand   <- add_len(Positive_strand[Positive_strand$Portion=="exon",])

  Positive_strand_transcript <- add_len(Positive_strand_transcript)
  Positive_strand_start_codon <- add_len(Positive_strand_start_codon)
  Positive_strand_stop_codon <- add_len(Positive_strand_stop_codon)
  five_prime_UTR_positive <- add_len(five_prime_UTR_positive)
  Three_prime_UTR_positive <- add_len(Three_prime_UTR_positive)

  Positive_strand_transcript$protein_sequence <-
    safe_protein_map(
      Positive_strand_transcript$Reference_id,
      protein_dict
    )

  ## ==================== NEGATIVE STRAND ====================
  Negative_strand_transcript <- Negative_strand[Negative_strand$Portion=="transcript",]
  Negative_strand_start_codon <- Negative_strand[Negative_strand$Portion=="start_codon",]
  Negative_strand_stop_codon  <- Negative_strand[Negative_strand$Portion=="stop_codon",]

  CDS_negative_strand    <- add_len(Negative_strand[Negative_strand$Portion=="CDS",])
  intron_negative_strand <- add_len(Negative_strand[Negative_strand$Portion=="intron",])
  exon_negative_strand   <- add_len(Negative_strand[Negative_strand$Portion=="exon",])

  ## ---- 5' UTR (Negative)
  Five_UTR_negative_strand <- cbind(
    Negative_strand_transcript[,!(names(Negative_strand_transcript) %in% c("Start","End"))],
    Start=Negative_strand_start_codon$End + 1,
    End=Negative_strand_transcript$End
  )
  Five_UTR_negative_strand$Portion <- "five_Prime_UTR"
  Five_UTR_negative_strand <- add_len(Five_UTR_negative_strand)

  ## ---- 3' UTR (Negative)
  Three_UTR_negative_strand <- cbind(
    Negative_strand_transcript[,!(names(Negative_strand_transcript) %in% c("Start","End"))],
    Start=Negative_strand_transcript$Start,
    End=Negative_strand_stop_codon$Start - 1
  )
  Three_UTR_negative_strand$Portion <- "three_Prime_UTR"
  Three_UTR_negative_strand <- add_len(Three_UTR_negative_strand)

  Negative_strand_transcript <- add_len(Negative_strand_transcript)
  Negative_strand_transcript$protein_sequence <-
    safe_protein_map(
      Negative_strand_transcript$Reference_id,
      protein_dict
    )

  ## ---- Promoter (Negative)
  Promoter_negative_strand <- cbind(
    Negative_strand_transcript[,!(names(Negative_strand_transcript) %in% c("Start","End","protein_sequence"))],
    Start=Negative_strand_transcript$End + 1,
    End=Negative_strand_transcript$End + PROMOTER_LEN
  )
  Promoter_negative_strand$Portion <- "Promoter"
  Promoter_negative_strand <- add_len(Promoter_negative_strand)

  ## ==================== Promoter (Positive) ====================
  promoter_positive <- data.frame(
    ID=five_prime_UTR_positive$ID,
    Source=five_prime_UTR_positive$Source,
    Portion="promoter",
    Start=five_prime_UTR_positive$Start - PROMOTER_LEN,
    End=five_prime_UTR_positive$Start - 1,
    score=five_prime_UTR_positive$score,
    Strand=five_prime_UTR_positive$Strand,
    Phase=".",
    Reference_id=five_prime_UTR_positive$Reference_id,
    stringsAsFactors = FALSE
  )
  promoter_positive <- add_len(promoter_positive)

  ## ==================== SAVE ALL (18 FILES) ====================
  outputs <- list(
    Positive_strand_transcript = Positive_strand_transcript,
    Positive_strand_start_codon = Positive_strand_start_codon,
    CDS_Positive_strand = CDS_Positive_strand,
    intron_positive_strand = intron_positive_strand,
    exon_positive_strand = exon_positive_strand,
    five_prime_UTR_positive = five_prime_UTR_positive,
    Three_prime_UTR_positive = Three_prime_UTR_positive,
    Positive_strand_stop_codon = Positive_strand_stop_codon,
    promoter_positive = promoter_positive,

    Negative_strand_transcript = Negative_strand_transcript,
    Negative_strand_start_codon = Negative_strand_start_codon,
    Negative_strand_stop_codon = Negative_strand_stop_codon,
    CDS_negative_strand = CDS_negative_strand,
    intron_negative_strand = intron_negative_strand,
    exon_negative_strand = exon_negative_strand,
    Five_UTR_negative_strand = Five_UTR_negative_strand,
    Three_UTR_negative_strand = Three_UTR_negative_strand,
    Promoter_negative_strand = Promoter_negative_strand
  )

  for (nm in names(outputs)) {
    write.csv(
      outputs[[nm]],
      file.path(output_dir, paste0(nm, ".csv")),
      row.names = FALSE
    )
  }

  message("ALL 18 FILES GENERATED SUCCESSFULLY (NO ERRORS)")
}



