"""MSSQL (SQL Server) schema introspection using pymssql + asyncio.to_thread."""
from __future__ import annotations

import asyncio
from urllib.parse import urlparse

from ._schema import ColumnInfo, DBSchema, TableSchema


def _parse_mssql_url(conn_str: str) -> dict:
    p = urlparse(conn_str)
    return {
        "server": p.hostname or "localhost",
        "port": p.port or 1433,
        "user": p.username or "",
        "password": p.password or "",
        "database": (p.path or "").lstrip("/"),
    }


def _load_mssql_schema_sync(conn_str: str) -> DBSchema:
    import pymssql
    params = _parse_mssql_url(conn_str)
    conn = pymssql.connect(**params)
    tables: dict[str, TableSchema] = {}
    try:
        cur = conn.cursor()
        cur.execute(
            "SELECT TABLE_NAME FROM INFORMATION_SCHEMA.TABLES "
            "WHERE TABLE_TYPE='BASE TABLE' ORDER BY TABLE_NAME"
        )
        table_names = [r[0] for r in cur.fetchall()]

        for tname in table_names:
            cur.execute(
                "SELECT COLUMN_NAME, DATA_TYPE, IS_NULLABLE "
                "FROM INFORMATION_SCHEMA.COLUMNS "
                "WHERE TABLE_NAME=%s ORDER BY ORDINAL_POSITION",
                (tname,),
            )
            cols_raw = cur.fetchall()

            cur.execute(
                "SELECT ku.COLUMN_NAME "
                "FROM INFORMATION_SCHEMA.TABLE_CONSTRAINTS tc "
                "JOIN INFORMATION_SCHEMA.KEY_COLUMN_USAGE ku "
                "  ON tc.CONSTRAINT_NAME=ku.CONSTRAINT_NAME "
                "WHERE tc.CONSTRAINT_TYPE='PRIMARY KEY' AND tc.TABLE_NAME=%s",
                (tname,),
            )
            pks = {r[0] for r in cur.fetchall()}
            columns = [
                ColumnInfo(name=r[0], type=r[1], not_null=(r[2] == "NO"), pk=(r[0] in pks))
                for r in cols_raw
            ]

            cur.execute(f"SELECT COUNT(*) FROM [{tname}]")
            row_count = cur.fetchone()[0]

            cur.execute(f"SELECT TOP 3 * FROM [{tname}]")
            raw = cur.fetchall()
            col_names = [d[0] for d in cur.description]
            sample = [dict(zip(col_names, r)) for r in raw]

            tables[tname] = TableSchema(
                name=tname, columns=columns,
                sample_rows=sample, row_count=row_count, foreign_keys=[],
            )
    finally:
        conn.close()
    return DBSchema(tables=tables, db_type="mssql")


async def load_mssql_schema(conn_str: str) -> DBSchema:
    return await asyncio.to_thread(_load_mssql_schema_sync, conn_str)