"""
FLYNC Bioinformatics Pipeline - Main Snakefile

This Snakefile orchestrates the RNA-seq analysis pipeline for lncRNA discovery.
It includes read mapping, transcriptome assembly, merging, and quantification.
"""

import pandas as pd
from pathlib import Path
import os
import re

# Note: Configuration is loaded via --configfile CLI argument
# No hardcoded configfile directive to avoid loading defaults

# Parse sample metadata - resolve path relative to working directory or as absolute
samples_path_config = config.get("samples", None)

# Auto-detect samples from FASTQ directory if no sample file provided
if samples_path_config is None and config.get("fastq_dir", None) is not None:
    # Auto-detect mode: scan FASTQ directory for sample names
    fastq_dir = Path(config["fastq_dir"])
    if not fastq_dir.exists():
        raise FileNotFoundError(f"FASTQ directory not found: {fastq_dir}")
    
    # Scan for FASTQ files
    import re
    sample_set = set()
    
    # Look for paired-end patterns: sample_1.fastq.gz, sample_2.fastq.gz
    for pattern in ["*_1.fastq.gz", "*_1.fq.gz", "*.fastq.gz", "*.fq.gz"]:
        for fq_file in fastq_dir.glob(pattern):
            # Remove suffixes to get sample name
            sample_name = fq_file.stem.replace('.fastq', '').replace('.fq', '')
            # Remove _1, _2, _R1, _R2 suffixes for paired-end
            sample_name = re.sub(r'[_\.][12]$', '', sample_name)
            sample_name = re.sub(r'[_\.]R[12]$', '', sample_name, flags=re.IGNORECASE)
            sample_set.add(sample_name)
    
    SAMPLES = sorted(list(sample_set))
    
    if len(SAMPLES) == 0:
        raise ValueError(
            f"No FASTQ files found in {fastq_dir}\n"
            f"  Expected files like: sample_1.fastq.gz, sample_2.fastq.gz\n"
            f"  Or: sample.fastq.gz"
        )
    
    print(f"Auto-detected {len(SAMPLES)} samples from FASTQ directory: {SAMPLES}")

elif samples_path_config is not None:
    # Manual mode: read from specified file
    samples_path = Path(samples_path_config)
    if not samples_path.is_absolute():
        # Try relative to current working directory first
        if not samples_path.exists():
            # If not found, try relative to config file directory
            config_dir = Path(workflow.configfiles[0]).parent if workflow.configfiles else Path.cwd()
            samples_path = config_dir / samples_path
            
    if not samples_path.exists():
        raise FileNotFoundError(
            f"Sample metadata file not found: {samples_path_config}\n"
            f"  Looked in: {samples_path.resolve()}\n"
            f"  Working directory: {Path.cwd()}\n"
            f"  Please provide the correct path in your config file or omit 'samples' to auto-detect."
        )

    # Read sample list - support both CSV and plain text formats
    if samples_path.suffix in ['.csv', '.tsv']:
        # CSV/TSV format: prefer a header with a 'sample_id' column; gracefully fall back if absent
        def _looks_like_sample_id(x: str) -> bool:
            try:
                return bool(re.match(r'^(SRR|ERR|DRR)\d+$', str(x)))
            except Exception:
                return False

        # Use appropriate separator
        sep = '\t' if samples_path.suffix == '.tsv' else ','

        # First, read assuming header
        samples_df = pd.read_csv(samples_path, sep=sep)

        if 'sample_id' in samples_df.columns:
            # Standard metadata with header row
            SAMPLES = samples_df['sample_id'].astype(str).tolist()
        else:
            # Heuristics: if the first column name itself looks like a sample ID and
            # the first cell does not, treat file as header-less and re-read without header.
            first_col_name = str(samples_df.columns[0])
            first_cell = str(samples_df.iloc[0, 0]) if not samples_df.empty else ''
            if _looks_like_sample_id(first_col_name) and not _looks_like_sample_id(first_cell):
                samples_df = pd.read_csv(samples_path, sep=sep, header=None)
                SAMPLES = samples_df.iloc[:, 0].astype(str).tolist()
            else:
                # Assume first column holds sample IDs with a header (e.g., unnamed but not a sample ID)
                SAMPLES = samples_df.iloc[:, 0].astype(str).tolist()
    elif samples_path.suffix == '.txt':
        # Plain text format: one sample ID per line
        with open(samples_path) as f:
            SAMPLES = [line.strip() for line in f if line.strip() and not line.startswith('#')]
    else:
        # Try to read as CSV by default
        try:
            samples_df = pd.read_csv(samples_path)
            # Check if likely no header
            first_col_name = samples_df.columns[0]
            if first_col_name.startswith(('SRR', 'ERR', 'DRR')) or ' ' not in str(first_col_name):
                samples_df = pd.read_csv(samples_path, header=None)
            SAMPLES = samples_df.iloc[:, 0].tolist()
        except:
            # Fall back to plain text
            with open(samples_path) as f:
                SAMPLES = [line.strip() for line in f if line.strip() and not line.startswith('#')]
    
    print(f"Loaded {len(SAMPLES)} samples from {samples_path}: {SAMPLES}")

else:
    # SRA mode: must have sample file
    raise ValueError(
        "No samples specified. Either:\n"
        "  1. Set 'samples' in config.yaml to a CSV/TXT file with sample IDs, OR\n"
        "  2. Use --fastq-dir to auto-detect samples from FASTQ filenames"
    )

# Define output directory
OUTPUT_DIR = Path(config["output_dir"])

# Library layout configuration
# Supports 3 modes:
#   1. Global setting: fastq_paired = true/false (applies to all samples)
#   2. Mapping file: library_layout_file = path to CSV with sample-specific settings
#   3. Auto-detection: omit both for SRA auto-detect or local FASTQ pattern detection

LAYOUT_STORAGE_FILE = OUTPUT_DIR / ".library_layouts.json"
USE_SRA_MODE = config.get("fastq_dir", None) is None
USE_LAYOUT_MAPPING = config.get("library_layout_file", None) is not None

# Import layout manager
import sys
sys.path.insert(0, str(Path(workflow.basedir) / "scripts"))
from library_layout_manager import LibraryLayoutManager, load_layout_mapping

# Initialize layout manager
layout_manager = LibraryLayoutManager(LAYOUT_STORAGE_FILE)

if USE_LAYOUT_MAPPING:
    # Mode 2: Load from mapping file
    mapping_file = Path(config["library_layout_file"])
    if not mapping_file.is_absolute():
        if not mapping_file.exists() and workflow.configfiles:
            config_dir = Path(workflow.configfiles[0]).parent
            mapping_file = config_dir / mapping_file
    
    print(f"\nLoading per-sample library layouts from: {mapping_file}")
    try:
        layouts = load_layout_mapping(mapping_file)
        layout_manager.bulk_set_layouts(layouts)
        
        # Print summary
        paired_count = sum(1 for is_paired in layouts.values() if is_paired)
        single_count = len(layouts) - paired_count
        print(f"  ✓ Loaded {len(layouts)} samples: {paired_count} paired-end, {single_count} single-end")
        
        # Show first few samples
        for sample_id, is_paired in list(layouts.items())[:3]:
            print(f"    - {sample_id}: {'PAIRED' if is_paired else 'SINGLE'}")
        if len(layouts) > 3:
            print(f"    ... and {len(layouts) - 3} more")
    except Exception as e:
        raise ValueError(f"Failed to load library layout mapping: {e}")

elif "fastq_paired" in config:
    # Mode 1: Global setting
    is_paired = config["fastq_paired"]
    print(f"\nUsing global library layout: {'PAIRED-END' if is_paired else 'SINGLE-END'} for all {len(SAMPLES)} samples")
    
    # Store for all samples
    layouts = {sample: is_paired for sample in SAMPLES}
    layout_manager.bulk_set_layouts(layouts)

elif USE_SRA_MODE and len(SAMPLES) > 0:
    # Mode 3a: SRA auto-detection
    print("\nAuto-detecting library layouts from SRA metadata (per-sample)...")
    print("  NOTE: This queries NCBI for each sample. Set 'library_layout_file' to skip in future runs.")
    
    # We'll detect per-sample during download rule
    # For now, just inform the user
    print(f"  Will auto-detect layout for {len(SAMPLES)} samples during download")

elif not USE_SRA_MODE:
    # Mode 3b: Local FASTQ auto-detection
    print("\nWill auto-detect library layouts from local FASTQ file patterns")
    print(f"  Paired-end: *_1.fastq.gz + *_2.fastq.gz")
    print(f"  Single-end: *.fastq.gz (no _1/_2 suffix)")

# Include rule modules
include: "rules/mapping.smk"
include: "rules/assembly.smk"
include: "rules/merge.smk"
include: "rules/quantify.smk"
include: "rules/dge.smk"

# Define the final target rule
rule all:
    input:
        # Merged assembly
        OUTPUT_DIR / "assemblies/merged.gtf",
        OUTPUT_DIR / "assemblies/assembled-new-transcripts.fa",
        # Comparison results
        OUTPUT_DIR / "gffcompare/gffcmp.stats",
        # Quantification (if samples provided)
        OUTPUT_DIR / "cov/quantification_complete.txt" if len(SAMPLES) > 0 else [],
        # Differential expression (if metadata CSV provided with condition column)
        OUTPUT_DIR / "dge/dge_summary.csv" if RUN_DGE else []
