"""
Generate resolve data from dbpedia dumps

Data source:
https://databus.dbpedia.org/dbpedia/generic#files
"""

import bz2
import re
from typing import Generator

from anystore.io import logged_items, smart_open, smart_stream
from anystore.types import BytesGenerator, Uri
from anystore.util import get_extension
from pydantic import BaseModel

DBPEDIA_RE = re.compile(
    r'^<http:\/\/[a-z]+\.dbpedia.org\/resource\/(?P<id>.+)>\s<(?P<prop>.*)>\s"(?P<value>.*)"@(?P<lang>[a-z]+)\s\.'
)


class DbpediaRow(BaseModel):
    id: str
    prop: str
    value: str
    lang: str


def unpack_dbpedia_row(line: bytes) -> DbpediaRow | None:
    # <http://de.dbpedia.org/resource/Angela_Merkel> <http://xmlns.com/foaf/0.1/givenName> "Angela"@de .
    data = DBPEDIA_RE.match(line.decode())
    if data is not None:
        return DbpediaRow(**data.groupdict())


def stream_bz2(uri: Uri) -> BytesGenerator:
    with smart_open(uri) as fh:
        with bz2.open(fh) as f:
            yield from f.readlines()


def stream_dbpedia(uri: Uri) -> Generator[DbpediaRow, None, None]:
    if get_extension(uri) == "bz2":
        stream = stream_bz2(uri)
    else:
        stream = smart_stream(uri, "rb")
    for line in logged_items(stream, "Read", 100_000, "DBPedia row"):
        data = unpack_dbpedia_row(line)
        if data is not None:
            yield data
