import sqlite3

from jinja2 import Environment, FileSystemLoader


def get_tables_and_columns(conn):
    cursor = conn.cursor()
    cursor.execute("SELECT name FROM sqlite_master WHERE type='table';")
    tables = [
        {
            "name": table[0],
            "columns": get_columns(conn, table[0]),
            "foreign_keys": get_foreign_keys(conn, table[0]),
        }
        for table in cursor.fetchall()
    ]
    return tables


def get_columns(conn, table_name):
    cursor = conn.cursor()
    cursor.execute(f"PRAGMA table_info({table_name});")
    columns = [row[1] for row in cursor.fetchall()]
    return columns


def get_foreign_keys(conn, table_name):
    cursor = conn.cursor()
    cursor.execute(f"PRAGMA foreign_key_list({table_name});")
    foreign_keys = [
        {"id": row[0], "from": row[3], "to_table": row[2], "to": row[4]}
        for row in cursor.fetchall()
    ]
    return foreign_keys


def generate_links(tables):
    links = []
    for t_index, table in enumerate(tables):
        for fk in table["foreign_keys"]:
            target_index = next(
                i for i, target in enumerate(tables) if target["name"] == fk["to_table"]
            )
            source_y = 40 + table["columns"].index(fk["from"]) * 20
            target_y = 40 + tables[target_index]["columns"].index(fk["to"]) * 20
            links.append(
                {
                    "source": {"x": 50 + t_index * 150 + 120, "y": 50 + source_y},
                    "target": {"x": 50 + target_index * 150, "y": 50 + target_y},
                }
            )
    return links


def generate_er_diagram(database_path):
    conn = sqlite3.connect(database_path)
    tables = get_tables_and_columns(conn)
    links = []  # Currently, we won't extract relationships
    links = generate_links(tables)

    env = Environment(loader=FileSystemLoader("templates"))
    template = env.get_template("er_diagram.html")

    with open("index.html", "w") as f:
        f.write(template.render(tables=tables, links=links))


if __name__ == "__main__":
    db_path = "database.db"
    generate_er_diagram(db_path)