import cloudpickle
import fsspec
import pyarrow as pa
import pytest
from pytest_manifest import Manifest
from utils import Input

import oxbow.core as ox


class TestSamFile:
    @pytest.mark.parametrize(
        "filepath",
        ["data/sample.sam", "data/malformed.sam", "data/does-not-exist.sam"],
    )
    def test_init_callstack(self, filepath, wiretap, manifest: Manifest):
        with wiretap(ox.SamFile) as stack:
            try:
                ox.SamFile(filepath)
            except BaseException:
                pass
            finally:
                assert (
                    manifest[f"{ox.SamFile.__name__}({Input(filepath)})"]
                ) == "\n".join([c.serialize() for c in stack])

    @pytest.mark.parametrize(
        "regions",
        [("foo",), ("foo", "bar"), ("foo", "bar", "baz"), ("*",), None],
    )
    def test_fragments(self, regions):
        for filepath in ("data/sample.sam",):
            fragments = ox.SamFile(filepath, regions=regions).fragments()
            assert len(fragments) == (len(regions) if regions else 1)

    def test_serialized_fragments(self):
        fragments = ox.SamFile(
            lambda: fsspec.open("data/sample.sam.gz", mode="rb").open(),
            index=lambda: fsspec.open("data/sample.sam.gz.tbi", mode="rb").open(),
            compressed=True,
            regions=["chr1", "chr2"],
        ).fragments()

        fragments = cloudpickle.loads(cloudpickle.dumps(fragments))

        assert [f.count_rows() for f in fragments] == [2, 1]

    @pytest.mark.parametrize(
        "fields",
        [
            "*",
            ["qname", "rname", "mapq"],
            ["qname", "rname", "foo"],
        ],
    )
    def test_batches(self, fields, manifest: Manifest):
        batches = ox.SamFile("data/sample.sam", fields=fields).batches()
        try:
            actual = {
                f"batch-{i:02}": pa.record_batch(b).to_pydict()
                for i, b in enumerate(batches)
            }
        except (OSError, ValueError) as e:
            actual = str(e)

        assert manifest[f"fields={fields}"] == actual

    def test_input_encodings(self):
        file = ox.SamFile("data/sample.sam", compressed=False, batch_size=3)
        assert next((file.batches())).num_rows <= 3

        with pytest.raises(OSError):
            file = ox.SamFile("data/sample.sam", compressed=True, batch_size=3)
            next((file.batches()))

        file = ox.SamFile("data/sample.sam.gz", compressed=True, batch_size=3)
        assert next((file.batches())).num_rows <= 3

        with pytest.raises(FileNotFoundError):
            file = ox.SamFile("doesnotexist.sam", compressed=False, batch_size=3)
            next((file.batches()))

    @pytest.mark.parametrize(
        "regions",
        [
            ["chr1"],
            ["chr1:17-32"],
            ["chr1:17-32", "chr1:30-37"],
        ],
    )
    def test_input_with_regions(self, regions):
        file = ox.SamFile(
            "data/sample.sam.gz",
            compressed=True,
            index="data/sample.sam.gz.tbi",
            regions=regions,
        )
        file.pl()

        file = ox.SamFile(
            "data/sample.sam.gz",
            compressed=True,
            index=None,  # inferred from name
            regions=regions,
        )
        file.pl()


class TestBamFile:
    @pytest.mark.parametrize(
        "filepath",
        ["data/sample.bam", "data/malformed.bam", "data/does-not-exist.bam"],
    )
    def test_init_callstack(self, filepath, wiretap, manifest: Manifest):
        with wiretap(ox.BamFile) as stack:
            try:
                ox.BamFile(filepath, compressed=True)
            except BaseException:
                pass
            finally:
                assert (
                    manifest[
                        f"{ox.BamFile.__name__}({Input(filepath)}, compressed=True)"
                    ]
                ) == "\n".join([c.serialize() for c in stack])

    @pytest.mark.parametrize(
        "fields",
        [
            None,
            ["qname", "rname", "mapq"],
            ["qname", "rname", "foo"],
        ],
    )
    def test_batches(self, fields, manifest: Manifest):
        batches = ox.BamFile(
            "data/sample.bam", fields=fields, compressed=True
        ).batches()
        try:
            actual = {
                f"batch-{i:02}": pa.record_batch(b).to_pydict()
                for i, b in enumerate(batches)
            }
        except (OSError, ValueError) as e:
            actual = str(e)

        assert manifest[f"fields={fields}"] == actual

    @pytest.mark.parametrize(
        "regions",
        [["foo"], ["foo", "bar"], ["foo", "bar", "baz"], ["*"], None],
    )
    def test_fragments(self, regions):
        for filepath in ("data/sample.bam",):
            fragments = ox.BamFile(
                filepath, regions=regions, compressed=True
            ).fragments()
            assert len(fragments) == (len(regions) if regions else 1)

    def test_serialized_fragments(self):
        fragments = ox.BamFile(
            lambda: fsspec.open("data/sample.bam", mode="rb").open(),
            index=lambda: fsspec.open("data/sample.bam.bai", mode="rb").open(),
            compressed=True,
            regions=["chr1"],
        ).fragments()

        fragments = cloudpickle.loads(cloudpickle.dumps(fragments))

        assert [f.count_rows() for f in fragments] == [4]

    def test_malformed(self):
        with pytest.raises(OSError):
            next(ox.BamFile("data/malformed.bam", compressed=True).batches())

    def test_input_encodings(self):
        file = ox.BamFile("data/sample.bam", compressed=True, batch_size=3)
        assert next((file.batches())).num_rows <= 3

        with pytest.raises(BaseException):
            file = ox.BamFile("data/sample.bam", compressed=False, batch_size=3)
            next((file.batches()))

        file = ox.BamFile("data/sample.ubam", compressed=False, batch_size=3)
        assert next((file.batches())).num_rows <= 3

        with pytest.raises(BaseException):
            file = ox.BamFile("data/sample.ubam", compressed=True, batch_size=3)
            next((file.batches()))

        with pytest.raises(BaseException):
            file = ox.BamFile("doesnotexist.bam", compressed=False, batch_size=3)
            next((file.batches()))

    @pytest.mark.parametrize(
        "regions",
        [
            ["chr1"],
            ["chr1:17-32"],
            ["chr1:17-32", "chr1:30-37"],
        ],
    )
    def test_input_with_regions(self, regions):
        file = ox.BamFile(
            "data/sample.bam",
            compressed=True,
            index="data/sample.bam.bai",
            regions=regions,
        )
        file.pl()

        file = ox.BamFile(
            "data/sample.bam",
            compressed=True,
            index=None,  # inferred from name
            regions=regions,
        )
        file.pl()


class TestCramFile:
    @pytest.mark.parametrize(
        "filepath",
        ["data/sample.cram", "data/malformed.cram", "data/does-not-exist.cram"],
    )
    def test_init_callstack(self, filepath, wiretap, manifest: Manifest):
        with wiretap(ox.CramFile) as stack:
            try:
                ox.CramFile(filepath)
            except BaseException:
                pass
            finally:
                assert (
                    manifest[f"{ox.CramFile.__name__}({Input(filepath)})"]
                ) == "\n".join([c.serialize() for c in stack])

    @pytest.mark.parametrize(
        "fields",
        [
            None,
            ["qname", "rname", "mapq"],
            ["qname", "rname", "foo"],
        ],
    )
    def test_batches(self, fields, manifest: Manifest):
        batches = ox.CramFile(
            "data/sample.cram", fields=fields, compressed=True
        ).batches()
        try:
            actual = {
                f"batch-{i:02}": pa.record_batch(b).to_pydict()
                for i, b in enumerate(batches)
            }
        except (OSError, ValueError) as e:
            actual = str(e)

        assert manifest[f"fields={fields}"] == actual

    @pytest.mark.parametrize(
        "regions",
        [["foo"], ["foo", "bar"], ["foo", "bar", "baz"], ["*"], None],
    )
    def test_fragments(self, regions):
        for filepath in ("data/sample.cram",):
            fragments = ox.CramFile(
                filepath,
                regions=regions,
            ).fragments()
            assert len(fragments) == (len(regions) if regions else 1)

    def test_serialized_fragments(self):
        fragments = ox.CramFile(
            lambda: fsspec.open("data/sample.cram", mode="rb").open(),
            index=lambda: fsspec.open("data/sample.cram.crai", mode="rb").open(),
            regions=["chr1"],
        ).fragments()

        fragments1 = cloudpickle.loads(cloudpickle.dumps(fragments))

        assert [f.count_rows() for f in fragments1] == [2]

    @pytest.mark.parametrize(
        "regions",
        [
            ["chr1"],
            ["chr1:17-32"],
            ["chr1:17-32", "chr1:30-37"],
        ],
    )
    def test_input_with_regions(self, regions):
        file = ox.CramFile(
            "data/sample.cram",
            index="data/sample.cram.crai",
            regions=regions,
        )
        file.pl()

        file = ox.CramFile(
            "data/sample.cram",
            index=None,  # inferred from name
            regions=regions,
        )
        file.pl()

    def test_with_reference(self):
        file = ox.CramFile(
            "data/sample-ref.cram",
            reference="data/sample-ref.fa",
            reference_index="data/sample-ref.fa.fai",
        )
        batch = pa.record_batch(next(file.batches()))
        assert batch.num_rows == 5
        # Verify bases are resolved (not Ns)
        seqs = batch.column("seq").to_pylist()
        assert all("N" not in s for s in seqs if s is not None)

    def test_with_tags(self):
        file = ox.CramFile("data/sample.cram").with_tags()
        assert len(file.tag_defs) > 0
        batch = pa.record_batch(next(file.batches()))
        assert "tags" in batch.schema.names

    def test_with_tags_with_reference(self):
        # Verifies that _tag_discovery_kwargs passes reference to the scanner
        file = ox.CramFile(
            "data/sample-ref.cram",
            reference="data/sample-ref.fa",
            reference_index="data/sample-ref.fa.fai",
        ).with_tags()
        assert isinstance(file.tag_defs, list)

    def test_with_reference_and_regions(self):
        file = ox.CramFile(
            "data/sample-ref.cram",
            reference="data/sample-ref.fa",
            reference_index="data/sample-ref.fa.fai",
            index="data/sample-ref.cram.crai",
            regions=["chr1"],
        )
        batch = pa.record_batch(next(file.batches()))
        assert batch.num_rows == 3
        rnames = batch.column("rname").to_pylist()
        assert all(r == "chr1" for r in rnames)
