import csv
import os
from collections import defaultdict
from io import TextIOBase
from pathlib import Path
from typing import Any, Callable

from biobit.analysis import seqproj

RowHook = Callable[[dict[str, Any]], dict[str, Any] | None]


def _reroot_fastq(row: dict[str, Any], fastq_root: os.PathLike[str] | None) -> dict[str, Any]:
    if fastq_root is None:
        return row

    fastq_root = Path(fastq_root)
    for field in ['fastq_1', 'fastq_2']:
        if field in row:
            fname = Path(row[field]).name
            row[field] = (fastq_root / fname).as_posix()
    return row


def load_bioproj(
        samplesheet: os.PathLike[str] | TextIOBase, *,
        hook: RowHook | None = None, fastq_root: os.PathLike[str] | None = None
) -> tuple[seqproj.Project, ...]:
    """
    Load a samplesheet.csv generated by nf-core/fetchngs and return a tuple of parsed bioproj.Project objects.

    :param samplesheet: The path to the samplesheet.csv file or a file-like object.
    :param hook: An optional hook function that is called for each row in the samplesheet. This might be used to filter
        out rows or modify them before parsing (e.g., correct fastq paths).
    :param fastq_root: An optional root path for fastq file names. If provided, the fastq file **names**
        in the samplesheet will be prefixed with this path before parsing. Note that only file names are used,
        the rest of the path from the samplesheet is discarded (if any).
    :return: A tuple of parsed bioproj.Project objects.
    """
    # SeqRun
    # run_accession (SRR*) -> SeqRun ind
    # instrument_model, instrument_platform -> SeqRun machine
    # library_layout -> SeqRun layout
    # read_count -> SeqRun reads
    # base_count -> SeqRun bases
    # fastq_1, fastq_2 -> SeqRun files

    # Library
    # library_source -> Library source
    # library_selection -> Library selection
    # ? -> Library stranding (always unknown)

    # Sample
    # sample_accession (SAMN*) -> Sample ind
    # scientific_name -> Sample organism
    # sample_title -> Sample tags
    # sample_description -> Sample description

    # Experiment
    # experiment_accession (SRX*) -> Experiment ind
    # experiment_title -> Experiment description

    # Project
    # study_accession (PRJNA*) -> Project ind
    # study_title -> Project description

    stream: TextIOBase
    if isinstance(samplesheet, TextIOBase):
        stream, close_later = samplesheet, False
    else:
        stream, close_later = open(samplesheet, 'r'), True

    hook = hook or (lambda x: x)

    projects: dict[str, Any] = defaultdict(lambda *_: {"samples": [], "experiments": {}, "description": None})
    for row in csv.DictReader(stream):
        # Skip empty rows
        if not row or (hooked := hook(row)) is None:
            continue

        row = _reroot_fastq(hooked, fastq_root)

        # Check for missing fields
        for field in [
            # SeqRun fields
            'run_accession', 'instrument_model', 'instrument_platform', 'library_layout', 'read_count', 'base_count',
            'fastq_1',
            # Library fields
            'library_source', 'library_selection',
            # Sample fields
            'sample_accession', 'scientific_name', 'sample_title', 'sample_description',
            # Experiment fields
            'experiment_accession', 'experiment_title',
            # Project fields
            'study_accession', 'study_title'
        ]:
            if field not in row or not row[field]:
                raise ValueError(f"Missing field {field} in the fetch-ngs samplesheet, row: {row}")
        if row['library_layout'] == 'PAIRED' and 'fastq_2' not in row:
            raise ValueError(f"Missing field fastq_2 in the fetch-ngs samplesheet for paired-end run, row: {row}")

        prjind = row['study_accession']
        project = projects[prjind]

        # Add project description
        if project["description"] is None:
            project["description"] = row['study_title']
        elif project["description"] != row['study_title']:
            raise ValueError(
                f"Project description mismatch for project {prjind}. "
                f"Expected {project['description']} but got {row['study_title']}"
            )

        # Parse attributes
        attributes = {"title": row['sample_title']}

        if row['sample_description']:
            description = row['sample_description']
        else:
            description = None

        sample = seqproj.Sample(
            ind=row['sample_accession'],
            organism=(row['scientific_name'],),
            attributes=attributes,
            description=description
        )
        project["samples"].append(sample)

        # Parse the experiment run
        library = seqproj.Library(
            source=(row['library_source'],),
            selection=(row['library_selection'],),
            stranding=seqproj.Stranding.Unknown
        )
        files: tuple[Path, ...]
        match row['library_layout']:
            case 'PAIRED':
                layout = seqproj.SeqLayout.Paired
                files = (Path(row['fastq_1']), Path(row['fastq_2']))
            case 'SINGLE':
                layout = seqproj.SeqLayout.Single
                files = (Path(row['fastq_1']),)
            case _:
                raise ValueError(
                    f"Unknown library layout {row['library_layout']} in the fetch-ngs samplesheet, row: {row}"
                )

        read_count = int(row['read_count']) if 'read_count' in row else None
        base_count = int(row['base_count']) if 'base_count' in row else None
        seqrun = seqproj.SeqRun(
            ind=row['run_accession'],
            machine=f"{row['instrument_platform']}: {row['instrument_model']}",
            layout=layout,
            files=files,
            reads=read_count,
            bases=base_count
        )

        expind = row['experiment_accession']
        if expind not in project["experiments"]:
            project["experiments"][expind] = {
                "sample": sample, "library": library, "seqruns": [seqrun], "description": row['experiment_title']
            }
        else:
            experiment = project["experiments"][expind]
            if sample != experiment["sample"]:
                raise ValueError(
                    f"Sample metadata mismatch for experiment {expind}. "
                    f"Expected {experiment['sample']} but got {sample}"
                )
            if library != experiment["library"]:
                raise ValueError(
                    f"Library metadata mismatch for experiment {expind}. "
                    f"Expected {experiment['library']} but got {library}"
                )
            if row['experiment_title'] != experiment["description"]:
                raise ValueError(
                    f"Experiment description mismatch for experiment {expind}. "
                    f"Expected {experiment['description']} but got {row['experiment_title']}"
                )
            experiment["seqruns"].append(seqrun)

    result = []
    for prjind, data in projects.items():
        experiments, samples, description = data.pop("experiments"), data.pop("samples"), data.pop("description")

        # Create bioproj.Experiments
        experiments = tuple(
            seqproj.Experiment(
                ind=expind, sample=experiment["sample"], library=experiment["library"], runs=experiment["seqruns"],
                description=experiment["description"]
            )
            for expind, experiment in experiments.items()
        )

        result.append(seqproj.Project(prjind, experiments, samples, description))

    if close_later:
        stream.close()

    return tuple(result)
