# src/gene_fetch/core.py
"""
Core functionality for Gene Fetch.
Contains Config class and basic utility functions.
"""

import logging
import sys
from dataclasses import dataclass
from pathlib import Path
from typing import Optional, Set, Dict, List, Any, FrozenSet

logger = logging.getLogger("gene_fetch")

def make_out_dir(path: Path) -> None:
    """Ensure output directory exists and create if it doesn't."""
    path.mkdir(parents=True, exist_ok=True)

def setup_logging(output_dir: Path) -> logging.Logger:
    """Setup logging with both file and console handlers."""
    make_out_dir(output_dir)
    
    # Clear existing handlers
    logger.handlers.clear()
    logger.setLevel(logging.INFO)
    
    formatter = logging.Formatter("%(asctime)s [%(levelname)s] %(message)s")
    
    console_handler = logging.StreamHandler(sys.stdout)
    console_handler.setFormatter(formatter)
    logger.addHandler(console_handler)
    
    log_path = output_dir / "gene_fetch.log"
    file_handler = logging.FileHandler(log_path)
    file_handler.setFormatter(formatter)
    logger.addHandler(file_handler)
    
    logger.info(f"Logging initialised. Log file: {log_path}")
    return logger

def log_progress(current: int, total: int, interval: int = 10) -> None:
    """Log progress at specified intervals."""
    if current == 0:
        logger.info("")
        logger.info(f"Starting processing: 0/{total} samples processed (0%)")
    elif current == total:
        logger.info(f"Processed: {total}/{total} samples processed (100%)")
    elif current % interval == 0:
        percentage = (current / total) * 100
        logger.info(
            f"=====   Progress: {current}/{total} samples processed "
            f"({percentage:.2f}%)"
        )

def get_process_id_column(header):
    """Identify the ID column from possible variations."""
    # Print what's in the header
    logger.info("")
    logger.info(f"CSV header detected: {header}")
    
    valid_names = [
        "ID",
        "process_id",
        "Process ID",
        "process id",
        "Process id",
        "PROCESS ID",
        "sample",
        "SAMPLE",
        "Sample",
    ]
    
    # Print repr of each header item to see any hidden chars
    for i, col in enumerate(header):
        logger.info(f"Header column {i}: {repr(col)}")
    
    # Try direct comparison first
    for col in header:
        if col in valid_names:
            logger.info(f"Found matching column: {col}")
            return col
    
    # Try trimming whitespace (in case of spaces)
    for col in header:
        trimmed = col.strip()
        if trimmed in valid_names:
            logger.info(f"Found matching column after trimming: {trimmed}")
            return col
    
    # Try case-insensitive comparison as last attempt
    for col in header:
        if col.upper() in [name.upper() for name in valid_names]:
            logger.info(f"Found matching column case-insensitive: {col}")
            return col
    
    logger.error(f"No matching column found in {header}")
    return None


# =============================================================================
# Configuration
# =============================================================================
@dataclass
class Config:
    """Configuration class for Gene Fetch, containing all runtime settings."""
    def __init__(self, email, api_key):
        # Email and API key required
        if not email:
            raise ValueError(
                "Email address required for NCBI API requests. Use -e/--email "
                "to provide your email."
            )
        if not api_key:
            raise ValueError(
                "API key required for NCBI API requests. Use -k/--api-key "
                "to provide your API key."
            )


        # Set run parameters
        self.email = email
        self.api_key = api_key

        # With an API key, 10 requests per second can be made
        self.max_calls_per_second = 10

        # Default batch size for fetching sequences
        self.fetch_batch_size = 200

        # Delay between batches (seconds) - uniform random delay between 1-2 seconds
        self.batch_delay = (1, 2)

        # Set search 'type'
        self.valid_sequence_types = frozenset({"protein", "nucleotide", "both"})

        # Minimum nucleotide and protein lengths for 'batch' mode
        self.protein_length_threshold = 500
        self.nucleotide_length_threshold = 1000

        # Minimum nucleotide and protein lengths for 'single' mode
        self.min_nucleotide_size_single_mode = 200
        self.min_protein_size_single_mode = 100

        self.gene_search_term = ""
    
        # Define gene type categories
        self._rRNA_genes = {
            "16s": [
                "16S ribosomal RNA[Gene]",
                "16S rRNA[Gene]",
                "rrs[Gene]",
                "rrn16[Gene]",
                "16S ribosomal RNA[Title] NOT methylase[Title] NOT "
                "methyltransferase[Title] NOT pseudouridylate[Title] NOT "
                "synthase[Title]",
                "16S rRNA[Title] NOT methylase[Title] NOT "
                "methyltransferase[Title] NOT pseudouridylate[Title] NOT "
                "synthase[Title]",
                "16S[Title] NOT methylase[Title] NOT methyltransferase[Title] "
                "NOT pseudouridylate[Title] NOT synthase[Title]",
                "16S ribosomal RNA[All Fields]",
                "rrs[All Fields]",
                "rrn16[All Fields]",
            ],
            "18s": [
                "18S ribosomal RNA[Gene]",
                "18S rRNA[Gene]",
                "rrn18[Gene]",
                "SSU rRNA[Gene]",
                "18S ribosomal RNA[Title] NOT methylase[Title] NOT "
                "methyltransferase[Title] NOT pseudouridylate[Title] NOT "
                "synthase[Title]",
                "18S rRNA[Title] NOT methylase[Title] NOT "
                "methyltransferase[Title] NOT pseudouridylate[Title] NOT "
                "synthase[Title]",
                "18S[Title] NOT methylase[Title] NOT methyltransferase[Title] "
                "NOT pseudouridylate[Title] NOT synthase[Title]",
                "18S ribosomal RNA[rRNA]",
                "SSU rRNA[rRNA]",
                "rrn18[rRNA]",
            ],
            "23s": [
                "23S ribosomal RNA[Gene]",
                "23S rRNA[Gene]",
                "rrl[Gene]",
                "rrn23[Gene]",
                "23S ribosomal RNA[Title] NOT methylase[Title] NOT "
                "methyltransferase[Title] NOT pseudouridylate[Title] NOT "
                "synthase[Title]",
                "23S rRNA[Title] NOT methylase[Title] NOT "
                "methyltransferase[Title] NOT pseudouridylate[Title] NOT "
                "synthase[Title]",
                "23S[Title] NOT methylase[Title] NOT methyltransferase[Title] "
                "NOT pseudouridylate[Title] NOT synthase[Title]",
                "23S ribosomal RNA[rRNA]",
                "rrl[rRNA]",
                "rrn23[rRNA]",
            ],
            "28s": [
                "28S ribosomal RNA[Gene]",
                "28S rRNA[Gene]",
                "rrn28[Gene]",
                "LSU rRNA[Gene]",
                "28S ribosomal RNA[Title] NOT methylase[Title] NOT "
                "methyltransferase[Title] NOT pseudouridylate[Title] NOT "
                "synthase[Title]",
                "28S rRNA[Title] NOT methylase[Title] NOT "
                "methyltransferase[Title] NOT pseudouridylate[Title] NOT "
                "synthase[Title]",
                "28S[Title] NOT methylase[Title] NOT methyltransferase[Title] "
                "NOT pseudouridylate[Title] NOT synthase[Title]",
                "28S ribosomal RNA[rRNA]",
                "LSU rRNA[rRNA]",
                "rrn28[rRNA]",
            ],
            "12s": [
                "12S ribosomal RNA[Gene]",
                "12S rRNA[Gene]",
                "mt-rrn1[Gene]",
                "mt 12S rRNA[Gene]",
                "12S ribosomal RNA[Title] NOT methylase[Title] NOT "
                "methyltransferase[Title] NOT pseudouridylate[Title] NOT "
                "synthase[Title]",
                "12S rRNA[Title] NOT methylase[Title] NOT "
                "methyltransferase[Title] NOT pseudouridylate[Title] NOT "
                "synthase[Title]",
                "12S[Title] NOT methylase[Title] NOT methyltransferase[Title] "
                "NOT pseudouridylate[Title] NOT synthase[Title]",
                "12S ribosomal RNA[rRNA]",
                "mt-rrn1[rRNA]",
                "mt 12S rRNA[rRNA]",
            ],
            "its1": [
                "ITS1[Title]",
                "internal transcribed spacer 1[Title]",
                "ITS-1[Title]",
                "ITS 1[Title]",
                "ITS1 ribosomal DNA[Title]",
                "ITS1 rDNA[Title]",
            ],
            "its2": [
                "ITS2[Title]",
                "internal transcribed spacer 2[Title]",
                "ITS-2[Title]",
                "ITS 2[Title]",
                "ITS2 ribosomal DNA[Title]",
                "ITS2 rDNA[Title]",
            ],
            "its": [
                "ITS[Title]",
                "internal transcribed spacer[Title]",
                "ITS region[Title]",
                "ITS1-5.8S-ITS2[Title]",
                "ribosomal ITS[Title]",
                "rDNA ITS[Title]",
            ],
            "trnl": [
                "trnL[Title]",
                "trnL gene[Title]",
                "tRNA-Leu[Title]",
                "tRNA-Leucine[Title]",
                "trnL-trnF[Title]",
                "trnL-F[Title]",
                "chloroplast trnL[Title]",
                "cp trnL[Title]",
                "trnL intron[Title]",
                "trnL-UAA[Title]",
                '"trnL gene"[Gene]',
            ],
        }

        self._protein_coding_genes = {
            "cox1": [
                "cox1[Gene]",
                "COI[Gene]",
                '"cytochrome c oxidase subunit 1"[Protein Name]',
                '"cytochrome oxidase subunit 1"[Protein Name]',
                '"cytochrome c oxidase subunit I"[Protein Name]',
                '"COX1"[Protein Name]',
                '"COXI"[Protein Name]',
            ],
            "cox2": [
                "cox2[Gene]",
                "COII[Gene]",
                '"cytochrome c oxidase subunit 2"[Protein Name]',
                '"cytochrome oxidase subunit 2"[Protein Name]',
                '"cytochrome c oxidase subunit II"[Protein Name]',
                '"COX2"[Protein Name]',
                '"COXII"[Protein Name]',
            ],
            "cox3": [
                "cox3[Gene]",
                "COIII[Gene]",
                '"cytochrome c oxidase subunit 3"[Protein Name]',
                '"cytochrome oxidase subunit 3"[Protein Name]',
                '"cytochrome c oxidase subunit III"[Protein Name]',
                '"COX3"[Protein Name]',
                '"COXIII"[Protein Name]',
            ],
            "cytb": [
                "cytb[Gene]",
                "cob[Gene]",
                '"cytochrome b"[Protein Name]',
                '"cytochrome b"[Gene]',
                '"CYTB"[Protein Name]',
            ],
            "nd1": [
                "nd1[Gene]",
                "NAD1[Gene]",
                '"NADH dehydrogenase subunit 1"[Protein Name]',
                '"ND1"[Protein Name]',
            ],
            "nd2": [
                "nd2[Gene]",
                "NAD2[Gene]",
                '"NADH dehydrogenase subunit 2"[Protein Name]',
                '"ND2"[Protein Name]',
            ],
            "rbcl": [
                "rbcL[Gene]",
                "RBCL[Gene]",
                '"ribulose-1,5-bisphosphate carboxylase/oxygenase large '
                'subunit"[Protein Name]',
                '"ribulose 1,5-bisphosphate carboxylase/oxygenase large '
                'subunit"[Protein Name]',
                '"ribulose bisphosphate carboxylase large chain"[Protein Name]',
                '"RuBisCO large subunit"[Protein Name]',
                '"ribulose-1,5-bisphosphate carboxylase/oxygenase small '
                'subunit"[Protein Name]',
                '"ribulose 1,5-bisphosphate carboxylase/oxygenase small '
                'subunit"[Protein Name]',
                '"ribulose-1,5-bisphosphate carboxylase/oxygenase small '
                'chain"[Protein Name]',
                '"RuBisCO small subunit"[Protein Name]',
                '"Ribulose bisphosphate carboxylase/oxygenase ' 'activase"[Gene]',
                '"rbcL gene"[Gene]',
                '"RBCL gene"[Gene]',
                "rbcL[Title]",
                "RBCL[Title]",
            ],
            "matk": [
                "matK[Gene]",
                "MATK[Gene]",
                '"maturase K"[Protein Name]',
                '"maturase K"[Gene]',
                '"maturase-K"[Protein Name]',
                '"maturase-K"[Gene]',
                '"Maturase K"[Protein Name]',
                '"Maturase K"[Gene]',
                '"matK gene"[Gene]',
                '"MATK gene"[Gene]',
                '"trnK-matK"[Gene]',
                '"maturase type II intron splicing factor"[Protein Name]',
                '"chloroplast group II intron splicing factor maturase '
                'K"[Protein Name]',
                '"type II intron maturase K"[Protein Name]',
                '"tRNA-lysine maturase K"[Protein Name]',
            ],
            "psba": [
                "psbA[Gene]",
                "PSBA[Gene]",
                '"photosystem II protein D1"[Protein Name]',
                '"photosystem II protein D1"[Gene]',
                '"PSII D1 protein"[Protein Name]',
                '"PSII D1 protein"[Gene]',
                '"photosystem II reaction center protein D1"[Protein Name]',
                '"photosystem II reaction center protein D1"[Gene]',
                '"photosystem Q(B) protein"[Protein Name]',
                '"32 kDa thylakoid membrane protein"[Protein Name]',
                '"psbA gene"[Gene]',
                '"PSBA gene"[Gene]',
                "psbA[Title]",
                "PSBA[Title]",
                '"trnH-psbA"[Title]',
                '"psbA-trnH"[Title]',
            ],
        }

    # Update sequence length thresholds
    def update_thresholds(self, protein_size: int, nucleotide_size: int):
        """Update the sequence length thresholds."""
        self.protein_length_threshold = protein_size
        self.nucleotide_length_threshold = nucleotide_size

    # Set search term based on gene name and type
    def set_gene_search_term(self, gene_name: str) -> str:
        """Set the gene search term based on the gene name and type."""
        gene_name = gene_name.lower()

        # Check if rRNA
        if gene_name in self._rRNA_genes:
            self.gene_search_term = "(" + " OR ".join(self._rRNA_genes[gene_name]) + ")"
            search_type = "rRNA"

        # Check if protein-coding
        elif gene_name in self._protein_coding_genes:
            self.gene_search_term = (
                "(" + " OR ".join(self._protein_coding_genes[gene_name]) + ")"
            )
            search_type = "protein-coding"

        else:
            # Generic search term for non-listed genes
            self.gene_search_term = (
                f"({gene_name}[Title] OR {gene_name}[Gene] OR "
                f'"{gene_name}"[Text Word])'
            )
            search_type = "generic"

        return search_type  # Return the type for logging purposes