import io
import codecs
import tarfile
import re
import gzip
import xml.etree.ElementTree as ET
from fnmatch import fnmatch
from pathlib import Path
from typing import NamedTuple
import ir_datasets
from ir_datasets.indices import PickleLz4FullStore
from .base import GenericDoc, GenericQuery, GenericScoredDoc, BaseDocs, BaseQueries, BaseScoredDocs, BaseQrels


class TrecDoc(NamedTuple):
    doc_id: str
    text: str
    marked_up_doc: str

class TitleUrlTextDoc(NamedTuple):
    doc_id: str
    title: str
    url: str
    text: str

class TrecQuery(NamedTuple):
    query_id: str
    title: str
    description: str
    narrative: str

class TrecSubtopic(NamedTuple):
    number: str
    text: str
    type: str

class TrecQrel(NamedTuple):
    query_id: str
    doc_id: str
    relevance: int
    iteration: str

class TrecPrel(NamedTuple):
    query_id: str
    doc_id: str
    relevance: int
    method: int
    iprob: float

# Default content tags from Anserini's TrecCollection
CONTENT_TAGS = 'TEXT HEADLINE TITLE HL HEAD TTL DD DATE LP LEADPARA'.split()

class TrecDocs(BaseDocs):
    def __init__(self, docs_dlc, encoding=None, path_globs=None, content_tags=CONTENT_TAGS, parser='BS4', namespace=None, lang=None, expected_file_count=None, docstore_size_hint=None, count_hint=None):
        self._docs_dlc = docs_dlc
        self._encoding = encoding
        self._path_globs = path_globs
        self._content_tags = content_tags
        self._parser = {
            'BS4': self._parser_bs,
            'text': self._parser_text,
            'tut': self._parser_tut,
        }[parser]
        self._doc = {
            'BS4': TrecDoc,
            'text': GenericDoc,
            'tut': TitleUrlTextDoc,
        }[parser]
        self._docs_namespace = namespace
        self._docs_lang = lang
        self._expected_file_count = expected_file_count
        self._docstore_size_hint = docstore_size_hint
        self._count_hint = count_hint
        if expected_file_count is not None:
            assert self._path_globs is not None, "expected_file_count only supported with path_globs"

    def docs_path(self, force=True):
        return self._docs_dlc.path(force)

    @ir_datasets.util.use_docstore
    def docs_iter(self):
        if Path(self._docs_dlc.path()).is_dir():
            if self._path_globs:
                file_count = 0
                for glob in sorted(self._path_globs):
                    for path in sorted(Path(self._docs_dlc.path()).glob(glob)):
                        file_count += 1
                        yield from self._docs_iter(path)
                if self._expected_file_count is not None:
                    if file_count != self._expected_file_count:
                        raise RuntimeError(f'found {file_count} files of the expected {self._expected_file_count} matching the following: {self._path_globs} under {self._docs_dlc.path()}. Make sure that directories are linked such that these globs match the correct number of files.')
            else:
                yield from self._docs_iter(self._docs_dlc.path())
        else:
            if self._path_globs:
                file_count = 0
                # tarfile, find globs, open in streaming mode (r|)
                with self._docs_dlc.stream() as stream:
                    with tarfile.open(fileobj=stream, mode='r|gz') as tarf:
                        for block in tarf:
                            if any(fnmatch(block.name, g) for g in self._path_globs):
                                file = tarf.extractfile(block)
                                if block.name.endswith('.gz'):
                                    file = gzip.GzipFile(fileobj=file)
                                yield from self._parser(file)
                                file_count += 1
                if self._expected_file_count is not None:
                    if file_count != self._expected_file_count:
                        raise RuntimeError(f'found {file_count} files of the expected {self._expected_file_count} matching the following: {self._path_globs} under {self._docs_dlc.path()}. Make sure that directories are linked such that these globs match the correct number of files.')
            else:
                with self._docs_dlc.stream() as f:
                    yield from self._parser(f)

    def _docs_iter(self, path):
        if Path(path).is_file():
            if str(path).endswith('.gz'):
                with gzip.open(path, 'rb') as f:
                    yield from self._parser(f)
            else:
                with path.open('rb') as f:
                    yield from self._parser(f)
        elif Path(path).is_dir():
            for child in path.iterdir():
                yield from self._docs_iter(child)

    def _parser_bs(self, stream):
        BeautifulSoup = ir_datasets.lazy_libs.bs4().BeautifulSoup
        f = codecs.getreader(self._encoding or 'utf8')(stream, errors='replace')
        doc_id, doc_markup = None, ''
        in_tag = False
        for line in f:
            if line.startswith('<DOCNO>'):
                doc_id = line.replace('<DOCNO>', '').replace('</DOCNO>\n', '').strip()
            elif line == '</DOC>\n':
                soup = BeautifulSoup(f'<OUTER>\n{doc_markup}\n</OUTER>', 'lxml')
                text = soup.get_text()
                yield TrecDoc(doc_id, text, doc_markup)
                doc_id, doc_markup = None, ''
            else:
                if in_tag:
                    doc_markup += line
                if line.startswith('</'):
                    if any(line.startswith(f'</{tag}>') for tag in self._content_tags):
                        in_tag -= 1
                if line.startswith('<'):
                    if any(line.startswith(f'<{tag}>') for tag in self._content_tags):
                        in_tag += 1
                        if in_tag == 1:
                            doc_markup += line

    def _parser_text(self, stream):
        f = codecs.getreader(self._encoding or 'utf8')(stream, errors='replace')
        doc_id, doc_text = None, ''
        in_tag = False
        for line in f:
            if line.startswith('<DOCNO>'):
                doc_id = line.replace('<DOCNO>', '').replace('</DOCNO>\n', '').strip()
            elif line == '</DOC>\n':
                yield GenericDoc(doc_id, doc_text)
                doc_id, doc_text = None, ''
            else:
                if line.startswith('</'):
                    if any(line.startswith(f'</{tag}>') for tag in self._content_tags):
                        in_tag = False
                if in_tag:
                    doc_text += line
                if line.startswith('<'):
                    if any(line.startswith(f'<{tag}>') for tag in self._content_tags):
                        in_tag = True

    def _parser_tut(self, stream):
        f = codecs.getreader(self._encoding or 'utf8')(stream, errors='replace')
        doc_id, doc_title, doc_url, doc_text = None, None, None, ''
        in_tag = False
        for line in f:
            if line.startswith('<DOCNO>'):
                doc_id = line.replace('<DOCNO>', '').replace('</DOCNO>\n', '').strip()
            if line.startswith('<TITLE>'):
                doc_title = line.replace('<TITLE>', '').replace('</TITLE>\n', '').strip()
            if line.startswith('<URL>'):
                doc_url = line.replace('<URL>', '').replace('</URL>\n', '').strip()
            elif line == '</DOC>\n':
                yield TitleUrlTextDoc(doc_id, doc_title, doc_url, doc_text)
                doc_id, doc_title, doc_url, doc_text = None, None, None, ''
            else:
                if line.startswith('</TEXT>'):
                    in_tag = False
                if in_tag:
                    doc_text += line
                if line.startswith('<TEXT>'):
                    in_tag = True

    def docs_cls(self):
        return self._doc

    def docs_store(self, field='doc_id'):
        return PickleLz4FullStore(
            path=f'{self.docs_path(force=False)}.pklz4',
            init_iter_fn=self.docs_iter,
            data_cls=self.docs_cls(),
            lookup_field=field,
            index_fields=['doc_id'],
            size_hint=self._docstore_size_hint,
            count_hint=self._count_hint,
        )

    def docs_count(self):
        if self.docs_store().built():
            return self.docs_store().count()

    def docs_namespace(self):
        return self._docs_namespace

    def docs_lang(self):
        return self._docs_lang


DEFAULT_QTYPE_MAP = {
    '<num> *(Number:)?': 'query_id',
    '<title> *(Topic:)?': 'title',
    '<desc> *(Description:)?': 'description',
    '<narr> *(Narrative:)?': 'narrative'
}
class TrecQueries(BaseQueries):
    def __init__(self, queries_dlc, qtype=TrecQuery, qtype_map=None, encoding=None, namespace=None, lang=None, remove_tags=('</title>',)):
        self._queries_dlc = queries_dlc
        self._qtype = qtype
        self._qtype_map = qtype_map or DEFAULT_QTYPE_MAP
        self._encoding = encoding
        self._queries_namespace = namespace
        self._queries_lang = lang
        self._remove_tags = remove_tags

    def queries_path(self):
        return self._queries_dlc.path()

    def queries_iter(self):
        fields, reading = {}, None
        with self._queries_dlc.stream() as f:
            f = codecs.getreader(self._encoding or 'utf8')(f)
            for line in f:
                if line.startswith('</top>'):
                    assert len(fields) == len(self._qtype._fields), fields
                    for tag in self._remove_tags:
                        fields = {k: v.replace(tag, '') for k, v in fields.items()}
                    yield self._qtype(*(fields[f].strip() for f in self._qtype._fields))
                    fields, reading = {}, None
                match_any = False
                for tag, target in self._qtype_map.items():
                    match = re.match(tag, line)
                    if match:
                        fields[target] = line[match.end():]
                        reading = target
                        match_any = True
                        break
                if not match_any and reading and not line.startswith('<'):
                    fields[reading] += line

    def queries_cls(self):
        return self._qtype

    def queries_namespace(self):
        return self._queries_namespace

    def queries_lang(self):
        return self._queries_lang


class TrecXmlQueries(BaseQueries):
    def __init__(self, queries_dlc, qtype=TrecQuery, qtype_map=None, encoding=None, subtopics_key='subtopics', namespace=None, lang=None):
        self._queries_dlc = queries_dlc
        self._qtype = qtype
        self._qtype_map = qtype_map or {f: f for f in qtype._fields}
        self._encoding = encoding
        self._subtopics_key = subtopics_key
        self._queries_namespace = namespace
        self._queries_lang = lang

    def queries_path(self):
        return self._queries_dlc.path()

    def queries_iter(self):
        with self._queries_dlc.stream() as f:
            f = codecs.getreader(self._encoding or 'utf8')(f)
            for topic_el in ET.fromstring(f.read()):
                item = [None for _ in self._qtype._fields]
                if 'number' in topic_el.attrib:
                    item[self._qtype._fields.index('query_id')] = topic_el.attrib['number']
                subtopics = []
                for attr in topic_el.attrib:
                    if attr in self._qtype_map:
                        text = topic_el.attrib[attr]
                        field = self._qtype_map[attr]
                        item[self._qtype._fields.index(field)] = text
                if topic_el.tag in self._qtype_map:
                    text = ''.join(topic_el.itertext())
                    field = self._qtype_map[topic_el.tag]
                    item[self._qtype._fields.index(field)] = text
                for field_el in topic_el:
                    if field_el.tag in self._qtype_map:
                        text = ''.join(field_el.itertext())
                        field = self._qtype_map[field_el.tag]
                        item[self._qtype._fields.index(field)] = text
                    if field_el.tag == 'subtopic':
                        text = ''.join(field_el.itertext())
                        subtopics.append(TrecSubtopic(field_el.attrib['number'], text, field_el.attrib['type']))
                if self._subtopics_key in self._qtype._fields:
                    item[self._qtype._fields.index('subtopics')] = tuple(subtopics)
                qid_field = self._qtype._fields.index('query_id')
                item[qid_field] = item[qid_field].strip() # remove whitespace from query_ids
                yield self._qtype(*item)

    def queries_cls(self):
        return self._qtype

    def queries_namespace(self):
        return self._queries_namespace

    def queries_lang(self):
        return self._queries_lang


class TrecColonQueries(BaseQueries):
    def __init__(self, queries_dlc, encoding=None, namespace=None, lang=None):
        self._queries_dlc = queries_dlc
        self._encoding = encoding
        self._queries_namespace = namespace
        self._queries_lang = lang

    def queries_iter(self):
        with self._queries_dlc.stream() as f:
            f = codecs.getreader(self._encoding or 'utf8')(f)
            for line in f:
                query_id, text = line.split(':', 1)
                text = text.rstrip('\n')
                yield GenericQuery(query_id, text)

    def queries_path(self):
        return self._queries_dlc.path()

    def queries_cls(self):
        return GenericQuery

    def queries_namespace(self):
        return self._queries_namespace

    def queries_lang(self):
        return self._queries_lang


class TrecQrels(BaseQrels):
    def __init__(self, qrels_dlc, qrels_defs):
        self._qrels_dlc = qrels_dlc
        self._qrels_defs = qrels_defs

    def qrels_path(self):
        return self._qrels_dlc.path()

    def qrels_iter(self):
        with self._qrels_dlc.stream() as f:
            f = codecs.getreader('utf8')(f)
            for line in f:
                if line == '\n':
                    continue # ignore blank lines
                cols = line.rstrip().split()
                if len(cols) != 4:
                    raise RuntimeError(f'expected 4 columns, got {len(cols)}')
                qid, it, did, score = cols
                yield TrecQrel(qid, did, int(score), it)

    def qrels_cls(self):
        return TrecQrel

    def qrels_defs(self):
        return self._qrels_defs


class TrecPrels(TrecQrels):
    def qrels_iter(self):
        with self._qrels_dlc.stream() as f:
            f = codecs.getreader('utf8')(f)
            for line in f:
                if line == '\n':
                    continue # ignore blank lines
                cols = line.rstrip().split()
                if len(cols) != 5:
                    raise RuntimeError(f'expected 5 columns, got {len(cols)}')
                qid, did, rel, method, iprob = cols
                yield TrecPrel(qid, did, int(rel), int(method), float(iprob))

    def qrels_cls(self):
        return TrecPrel


class TrecScoredDocs(BaseScoredDocs):
    def __init__(self, scoreddocs_dlc):
        self._scoreddocs_dlc = scoreddocs_dlc

    def scoreddocs_path(self):
        return self._scoreddocs_dlc.path()

    def scoreddocs_iter(self):
        with self._scoreddocs_dlc.stream() as f:
            f = codecs.getreader('utf8')(f)
            for line in f:
                cols = line.rstrip().split()
                if len(cols) == 6:
                    qid, _, did, _, score, _ = cols
                elif len(cols) == 2:
                    qid, did, score = *cols, '0'
                yield GenericScoredDoc(qid, did, float(score))

[(127, 'ir_datasets.lazy_libs.bs4', 'ir_datasets.lazy_libs.bs4', False, 'import ir_datasets\n'), (128, 'codecs.getreader', 'codecs.getreader', False, 'import codecs\n'), (152, 'codecs.getreader', 'codecs.getreader', False, 'import codecs\n'), (172, 'codecs.getreader', 'codecs.getreader', False, 'import codecs\n'), (115, 'pathlib.Path', 'Path', False, 'from pathlib import Path\n'), (240, 'codecs.getreader', 'codecs.getreader', False, 'import codecs\n'), (284, 'codecs.getreader', 'codecs.getreader', False, 'import codecs\n'), (332, 'codecs.getreader', 'codecs.getreader', False, 'import codecs\n'), (361, 'codecs.getreader', 'codecs.getreader', False, 'import codecs\n'), (381, 'codecs.getreader', 'codecs.getreader', False, 'import codecs\n'), (404, 'codecs.getreader', 'codecs.getreader', False, 'import codecs\n'), (117, 'gzip.open', 'gzip.open', False, 'import gzip\n'), (122, 'pathlib.Path', 'Path', False, 'from pathlib import Path\n'), (250, 're.match', 're.match', False, 'import re\n'), (99, 'tarfile.open', 'tarfile.open', False, 'import tarfile\n'), (101, 'fnmatch.fnmatch', 'fnmatch', False, 'from fnmatch import fnmatch\n'), (104, 'gzip.GzipFile', 'gzip.GzipFile', False, 'import gzip\n')]