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

import oxbow.core as ox
from tests.utils import Input


class TestBedFile:
    @pytest.mark.parametrize(
        "filepath",
        ["data/sample.bed", "data/malformed.bed", "data/does-not-exist.bed"],
    )
    def test_init_callstack(self, filepath, wiretap, manifest: Manifest):
        with wiretap(ox.BedFile) as stack:
            try:
                ox.BedFile(filepath)
            except BaseException:
                pass
            finally:
                assert (
                    manifest[f"{ox.BedFile.__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.bed",):
            fragments = ox.BedFile(filepath, regions=regions).fragments()
            assert len(fragments) == (len(regions) if regions else 1)

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

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

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

    @pytest.mark.parametrize(
        "fields",
        [
            None,
            "*",
            ("chrom", "start", "end"),
            ("nonexistent-field",),
        ],
    )
    def test_batches(self, fields, manifest: Manifest):
        batches = ox.BedFile("data/sample.bed", 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.BedFile("data/sample.bed", "bed9", compressed=False, batch_size=3)
        assert next((file.batches())).num_rows <= 3

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

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

        with pytest.raises(BaseException):
            file = ox.BedFile(
                "data/sample.bed.gz", "bed9", compressed=False, batch_size=3
            )
            next((file.batches()))

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

    @pytest.mark.parametrize(
        "regions",
        [
            ["chr1"],
            ["chr12"],
            ["chr5", "chr12:1-100"],
        ],
    )
    def test_input_with_regions(self, regions):
        file = ox.BedFile(
            "data/sample.bed.gz",
            compressed=True,
            index="data/sample.bed.gz.tbi",
            regions=regions,
        )
        file.pl()

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

    def test_projections(self):
        ds = ox.BedFile(
            "data/sample.bed", bed_schema="bed4", fields=["name", "end", "start"]
        )
        batch = next(ds.batches())
        assert list(batch.schema.names) == ["name", "end", "start"]
        with pytest.raises((OSError, ValueError)):
            next(
                ox.BedFile(
                    "data/sample.bed", bed_schema="bed4", fields=["rest"]
                ).batches()
            )

        ds = ox.BedFile(
            "data/sample.bed",
            bed_schema="bed4+2",
            fields=["BED4+2", "end", "BED4+1", "start"],
        )
        batch = next(ds.batches())
        # Extended fields get shuffled to the end in the order provided
        assert list(batch.schema.names) == ["end", "start", "BED4+2", "BED4+1"]
        with pytest.raises((OSError, ValueError)):
            next(
                ox.BedFile(
                    "data/sample.bed", bed_schema="bed4", fields=["BED4+3"]
                ).batches()
            )
        with pytest.raises((OSError, ValueError)):
            next(
                ox.BedFile(
                    "data/sample.bed", bed_schema="bed4", fields=["rest"]
                ).batches()
            )

        ds = ox.BedFile("data/sample.bed", bed_schema="bed4+", fields=["end", "start"])
        batch = next(ds.batches())
        assert list(batch.schema.names) == ["end", "start"]

        ds = ox.BedFile(
            "data/sample.bed", bed_schema="bed4+", fields=["end", "rest", "start"]
        )
        batch = next(ds.batches())
        # Extended fields get shuffled to the end
        assert list(batch.schema.names) == ["end", "start", "rest"]

    def test_custom_schema(self):
        import pyarrow as pa

        # sample.bed is a BED9 file. Reinterpret with custom column names
        # and types for the 6 extension columns beyond BED3.
        ds = ox.BedFile(
            "data/sample.bed",
            bed_schema=(
                "bed3",
                [
                    ("label", "string"),
                    ("score", "float"),
                    ("orientation", "string"),
                    ("block_start", "uint"),
                    ("block_end", "uint"),
                    ("color", "string"),
                ],
            ),
        )
        batch = next(ds.batches())
        assert list(batch.schema.names) == [
            "chrom",
            "start",
            "end",
            "label",
            "score",
            "orientation",
            "block_start",
            "block_end",
            "color",
        ]

        # Custom column types are respected
        assert batch.schema.field("label").type == pa.utf8()
        assert batch.schema.field("score").type == pa.float32()
        assert batch.schema.field("block_start").type == pa.uint32()

        # Dict form works the same way
        ds = ox.BedFile(
            "data/sample.bed",
            bed_schema=(
                "bed3",
                {"label": "string", "score": "float"},
            ),
        )
        batch = next(ds.batches())
        assert list(batch.schema.names) == [
            "chrom",
            "start",
            "end",
            "label",
            "score",
        ]

        # Custom schema with projection
        ds = ox.BedFile(
            "data/sample.bed",
            bed_schema=(
                "bed3",
                [("label", "string"), ("score_text", "string")],
            ),
            fields=["chrom", "label"],
        )
        batch = next(ds.batches())
        assert list(batch.schema.names) == ["chrom", "label"]

        # Projection of a non-first custom field (tests positional alignment)
        ds = ox.BedFile(
            "data/sample.bed",
            bed_schema=(
                "bed3",
                [("label", "string"), ("score_text", "string"), ("orient", "string")],
            ),
            fields=["chrom", "orient"],
        )
        batch = next(ds.batches())
        assert list(batch.schema.names) == ["chrom", "orient"]
        # orient is the 3rd custom field (6th overall), should have actual data
        assert batch.column("orient")[0].as_py() is not None
