from typing import AsyncIterable, Iterable

import pytest
from autobigs.engine.structures.alignment import AlignmentStats
from autobigs.engine.writing import alleles_to_text_map, write_mlst_profiles_as_csv
from autobigs.engine.structures.mlst import Allele, MLSTProfile, NamedMLSTProfile
import tempfile
from csv import reader
from os import path


@pytest.fixture
def dummy_alphabet_mlst_profile():
    return NamedMLSTProfile("name", MLSTProfile((
        Allele("A", "1", None),
        Allele("D", "1", None),
        Allele("B", "1", None),
        Allele("C", "1", None),
        Allele("C", "2", AlignmentStats(90, 10, 0, 90))
    ), "mysterious", "very mysterious"))

async def iterable_to_asynciterable(iterable: Iterable):
    for iterated in iterable:
        yield iterated

async def test_column_order_is_same_as_expected_file(dummy_alphabet_mlst_profile: MLSTProfile):
    dummy_profiles = [dummy_alphabet_mlst_profile]
    with tempfile.TemporaryDirectory() as temp_dir:
        output_path = path.join(temp_dir, "out.csv")
        await write_mlst_profiles_as_csv(iterable_to_asynciterable(dummy_profiles), output_path)
        with open(output_path) as csv_handle:
            csv_reader = reader(csv_handle)
            lines = list(csv_reader)
            target_columns = lines[0][3:]
            assert target_columns == sorted(target_columns)

async def test_csv_writing_sample_name_not_repeated_when_single_sequence(dummy_alphabet_mlst_profile):
    dummy_profiles = [dummy_alphabet_mlst_profile]
    with tempfile.TemporaryDirectory() as temp_dir:
        output_path = path.join(temp_dir, "out.csv")
        await write_mlst_profiles_as_csv(iterable_to_asynciterable(dummy_profiles), output_path)
        with open(output_path) as csv_handle:
            csv_reader = reader(csv_handle)
            lines = list(csv_reader)
            sample_name = lines[1][0]
            assert sample_name == "name"


async def test_alleles_to_text_map_mapping_is_correct(dummy_alphabet_mlst_profile: NamedMLSTProfile):
    assert dummy_alphabet_mlst_profile.mlst_profile is not None
    mapping = alleles_to_text_map(dummy_alphabet_mlst_profile.mlst_profile.alleles)
    expected_mapping = {
        "A": "1",
        "B": "1",
        "C": ("1", "2*"),
        "D": "1"
    }
    for allele_name, allele_ids in mapping.items():
        assert allele_name in expected_mapping
        assert allele_ids == expected_mapping[allele_name]

async def test_csv_writing_includes_asterisk_for_non_exact(dummy_alphabet_mlst_profile: NamedMLSTProfile):
    dummy_profiles = [dummy_alphabet_mlst_profile]
    with tempfile.TemporaryDirectory() as temp_dir:
        output_path = path.join(temp_dir, "out.csv")
        await write_mlst_profiles_as_csv(iterable_to_asynciterable(dummy_profiles), output_path)
        with open(output_path) as csv_handle:
            csv_reader = reader(csv_handle)
            lines = list(csv_reader)
            assert '*' in lines[1][5]