from contextlib import AbstractAsyncContextManager
import re
from typing import Sequence
from aiohttp import ClientSession, ClientTimeout
from mlstmyfasta.engine.data.MLST import Allele
from mlstmyfasta.engine.data.genomics import NamedString
from mlstmyfasta.engine.remote.databases.institutpasteur.structures import InstitutPasteurSequenceRequest, InstitutPasteurSequenceResponse

class InstitutPasteurProfiler(AbstractAsyncContextManager):

    async def __aenter__(self):
        return self


    def __init__(self, database_name: str):
        self._base_url = f"https://bigsdb.pasteur.fr/api/db/{database_name}/"
        self._http_client = ClientSession(self._base_url, timeout=ClientTimeout(10000))

    async def fetch_mlst_profile(self, sequence_string: str):
        # See https://bigsdb.pasteur.fr/api/db/pubmlst_bordetella_seqdef/schemes
        uri_path = f"schemes/3/sequence"
        response = await self._http_client.post(uri_path, json={
            "sequence": sequence_string
        })
        sequence_response: dict = await response.json()
        exact_matches: dict[str, Sequence[dict[str, str]]] = sequence_response["exact_matches"]  
        for allele_loci, alleles in exact_matches.items():
            for allele in alleles:
                alelle_id = allele["allele_id"]
                yield Allele(allele_loci=allele_loci, allele_variant=alelle_id)

    async def close(self):
        await self._http_client.close()

    async def __aexit__(self, exc_type, exc_value, traceback):
        await self.close()