#!/usr/bin/env python3

import argparse
import os
import sqlite3
import re


def parse_nt_line(line):
    """
    Parse simple N-Triples:
    <subject> <predicate> <object> .
    """
    line = line.rstrip("\n")

    if not line or line.startswith("#"):
        return None

    # split first two fields, keep object intact
    m = re.match(r'(<.*?>)\s+(<.*?>)\s+(.*?)\s+\.$', line)

    if not m:
        return None

    subject = m.group(1)[1:-1]
    predicate = m.group(2)[1:-1]
    obj = m.group(3)

    # remove literal quotes
    if obj.startswith('"'):
        obj = obj.split('"')[1]

    elif obj.startswith("<"):
        obj = obj[1:-1]

    return subject, predicate, obj


def sql_name(uri):
    """
    Convert RDF predicate URI into safe SQLite column name.
    """
    name = uri.split("/")[-1]
    name = name.split("#")[-1]

    name = re.sub(r'[^a-zA-Z0-9_]', '_', name)

    if not name:
        name = "column"

    return name


def discover_predicates(filename):
    predicates = {}

    print("Pass 1: discovering predicates...")

    with open(filename, "r", encoding="utf-8") as f:
        for line in f:
            triple = parse_nt_line(line)

            if triple:
                _, predicate, _ = triple

                if predicate not in predicates:
                    predicates[predicate] = sql_name(predicate)

    return predicates


def create_table(con, predicates):
    cur = con.cursor()

    columns = [
        "subject TEXT PRIMARY KEY"
    ]

    used = {"subject"}

    for col in predicates.values():

        original = col
        counter = 1

        while col in used:
            col = f"{original}_{counter}"
            counter += 1

        used.add(col)

        columns.append(f'"{col}" TEXT')

        predicates[
            next(k for k, v in predicates.items() if v == original)
        ] = col

    sql = """
    CREATE TABLE IF NOT EXISTS inspections (
        %s
    )
    """ % ",".join(columns)

    cur.execute(sql)
    con.commit()


def insert_row(cur, predicates, subject, values):

    cols = ["subject"]
    data = [subject]

    for pred, value in values.items():

        if pred in predicates:
            cols.append(predicates[pred])
            data.append(value)

    placeholders = ",".join("?" for _ in data)

    sql = (
        "INSERT INTO inspections "
        "(%s) VALUES (%s)"
        % (
            ",".join('"%s"' % c for c in cols),
            placeholders,
        )
    )

    cur.execute(sql, data)


def import_nt(filename, output):

    predicates = discover_predicates(filename)

    print(
        "Found %d predicates"
        % len(predicates)
    )

    con = sqlite3.connect(output)

    create_table(con, predicates)

    cur = con.cursor()

    print("Pass 2: importing triples...")

    current_subject = None
    row = {}

    count = 0

    with open(filename, "r", encoding="utf-8") as f:

        for line in f:

            triple = parse_nt_line(line)

            if not triple:
                continue

            subject, predicate, obj = triple

            if current_subject is None:
                current_subject = subject

            if subject != current_subject:

                insert_row(
                    cur,
                    predicates,
                    current_subject,
                    row
                )

                count += 1

                if count % 10000 == 0:
                    con.commit()
                    print(
                        "Inserted",
                        count,
                        "rows"
                    )

                current_subject = subject
                row = {}

            row[predicate] = obj


    if current_subject is not None:
        insert_row(
            cur,
            predicates,
            current_subject,
            row
        )

    con.commit()

    print(
        "Done. Inserted",
        count + 1,
        "rows"
    )

    con.close()


def main():

    parser = argparse.ArgumentParser()

    parser.add_argument(
        "-i",
        "--input",
        required=True
    )

    parser.add_argument(
        "-o",
        "--output",
        required=True
    )

    args = parser.parse_args()

    import_nt(
        args.input,
        args.output
    )


if __name__ == "__main__":
    main()