from datasette import hookimpl, Response
from datasette.resources import DatabaseResource
import re


def mermaid_quote(name):
    """Double-quote a name for mermaid ER diagrams.

    Handles reserved words (e.g. 'select', 'and'), names starting
    with digits, and special characters.
    """
    safe = re.sub(r"[^a-zA-Z0-9_ -]", "_", name)
    return f'"{safe}"'


MERMAID_RESERVED_WORDS = {"pk", "fk", "uk"}


def mermaid_word(name):
    """Sanitize a word for mermaid ER diagrams (unquoted).

    Appends underscore if name matches a reserved keyword (pk/fk/uk).
    """
    safe = re.sub(r"[^a-zA-Z0-9_]", "_", name)
    if safe.lower() in MERMAID_RESERVED_WORDS:
        safe = safe + "_"
    return safe


@hookimpl
def database_actions(datasette, actor, database):
    async def inner():
        if not await datasette.allowed(
            action="view-database",
            resource=DatabaseResource(database),
            actor=actor,
        ):
            return []
        return [
            {
                "href": datasette.urls.path(f"/-/table-diagram/{database}"),
                "label": "Table diagram",
                "description": "Entity Relationship diagram of tables",
            }
        ]

    return inner


@hookimpl
def register_routes():
    return [
        (r"^/-/table-diagram/(?P<database>.+)$", table_diagram_view),
    ]


async def table_diagram_view(request, datasette):
    database = request.url_vars["database"]
    try:
        db = datasette.get_database(database)
    except KeyError:
        return Response.html("Database not found", status=404)

    # Check view-database permission
    await datasette.ensure_permission(
        action="view-database",
        resource=DatabaseResource(database),
        actor=request.actor,
    )

    # Get tables the actor has view-table permission for
    visible_tables = set()
    page = await datasette.allowed_resources(
        "view-table", request.actor, parent=database
    )
    async for resource in page.all():
        visible_tables.add(resource.child)

    tables = [t for t in await db.table_names() if t in visible_tables]

    # Get foreign key relationships
    fk_result = await db.execute("""
        select
            sqlite_master.name as table_from,
            fk_info.[from] as column_from,
            fk_info.[table] as table_to,
            fk_info.[to] as column_to
        from
            sqlite_master
            join pragma_foreign_key_list(sqlite_master.name) as fk_info
        order by
            sqlite_master.name
        """)
    foreign_keys = [
        dict(zip(fk_result.columns, row))
        for row in fk_result.rows
        if row[0] in visible_tables and row[2] in visible_tables
    ]

    # Build set of FK columns for marking in diagram
    fk_columns = {(fk["table_from"], fk["column_from"]) for fk in foreign_keys}

    # Get columns for each table
    table_columns = {}
    for table in tables:
        col_result = await db.execute(
            "select * from pragma_table_info(:table)", {"table": table}
        )
        table_columns[table] = [
            dict(zip(col_result.columns, row)) for row in col_result.rows
        ]

    mermaid_text = generate_er_diagram(tables, foreign_keys, fk_columns, table_columns)

    return Response.html(
        await datasette.render_template(
            "table_diagram.html",
            {
                "database": database,
                "mermaid_text": mermaid_text,
                "tables": sorted(tables),
            },
            request=request,
        )
    )


def generate_er_diagram(tables, foreign_keys, fk_columns, table_columns):
    lines = ["erDiagram"]

    # Add relationships
    for fk in foreign_keys:
        table_to = mermaid_quote(fk["table_to"])
        table_from = mermaid_quote(fk["table_from"])
        column_from = mermaid_word(fk["column_from"])
        lines.append(f"    {table_to} ||--o{{ {table_from} : {column_from}")

    # Add table definitions with columns
    for table in sorted(tables):
        columns = table_columns.get(table, [])
        safe_table = mermaid_quote(table)
        lines.append(f"    {safe_table} {{")
        for col in columns:
            col_type = (col.get("type") or "text").split("(")[
                0
            ].strip().lower() or "text"
            col_type = mermaid_word(col_type)
            col_name = mermaid_word(col["name"])
            if col.get("pk"):
                marker = "PK"
            elif (table, col["name"]) in fk_columns:
                marker = "FK"
            else:
                marker = ""
            if marker:
                lines.append(f"        {col_type} {col_name} {marker}")
            else:
                lines.append(f"        {col_type} {col_name}")
        lines.append("    }")

    return "\n".join(lines)
