"""Schema command for DDL generation."""

import json
from dataclasses import dataclass


@dataclass
class ColumnDef:
    name: str
    data_type: str
    not_null: bool
    default: str | None
    identity: str  # '' | 'a' (always) | 'd' (by default)
    generated: str  # '' | 's' (stored)


@dataclass
class ConstraintDef:
    name: str
    type: str  # p=PK, u=unique, f=FK, c=check
    definition: str


@dataclass
class IndexDef:
    name: str
    definition: str


class SchemaError(Exception):
    """Raised when DDL cannot be generated."""

    pass


def table_exists(conn, schema: str, table: str) -> bool:
    """Check if table exists."""
    query = """
        SELECT 1 FROM pg_class c
        JOIN pg_namespace n ON n.oid = c.relnamespace
        WHERE n.nspname = %s AND c.relname = %s AND c.relkind = 'r'
    """
    with conn.cursor() as cur:
        cur.execute(query, (schema, table))
        return cur.fetchone() is not None


def get_columns(conn, schema: str, table: str) -> list[ColumnDef]:
    """Get column definitions using pg_attribute for accurate types."""
    query = """
        SELECT
            a.attname,
            format_type(a.atttypid, a.atttypmod),
            a.attnotnull,
            pg_get_expr(d.adbin, d.adrelid),
            a.attidentity,
            a.attgenerated
        FROM pg_attribute a
        LEFT JOIN pg_attrdef d ON d.adrelid = a.attrelid AND d.adnum = a.attnum
        JOIN pg_class c ON c.oid = a.attrelid
        JOIN pg_namespace n ON n.oid = c.relnamespace
        WHERE n.nspname = %s
          AND c.relname = %s
          AND a.attnum > 0
          AND NOT a.attisdropped
        ORDER BY a.attnum
    """
    with conn.cursor() as cur:
        cur.execute(query, (schema, table))
        rows = cur.fetchall()

    return [
        ColumnDef(
            name=r[0],
            data_type=r[1],
            not_null=r[2],
            default=r[3],
            identity=r[4] or "",
            generated=r[5] or "",
        )
        for r in rows
    ]


def get_constraints(conn, schema: str, table: str) -> list[ConstraintDef]:
    """Get constraint definitions (PK, UNIQUE, FK, CHECK)."""
    query = """
        SELECT
            c.conname,
            c.contype,
            pg_get_constraintdef(c.oid)
        FROM pg_constraint c
        JOIN pg_class t ON t.oid = c.conrelid
        JOIN pg_namespace n ON n.oid = t.relnamespace
        WHERE n.nspname = %s AND t.relname = %s
        ORDER BY
            CASE c.contype
                WHEN 'p' THEN 1
                WHEN 'u' THEN 2
                WHEN 'f' THEN 3
                WHEN 'c' THEN 4
            END,
            c.conname
    """
    with conn.cursor() as cur:
        cur.execute(query, (schema, table))
        rows = cur.fetchall()

    return [ConstraintDef(name=r[0], type=r[1], definition=r[2]) for r in rows]


def get_indexes(conn, schema: str, table: str) -> list[IndexDef]:
    """Get index definitions (excluding constraint-backing indexes)."""
    query = """
        SELECT
            i.relname,
            pg_get_indexdef(i.oid)
        FROM pg_index ix
        JOIN pg_class i ON i.oid = ix.indexrelid
        JOIN pg_class t ON t.oid = ix.indrelid
        JOIN pg_namespace n ON n.oid = t.relnamespace
        WHERE n.nspname = %s
          AND t.relname = %s
          AND i.oid NOT IN (
              SELECT conindid FROM pg_constraint
              WHERE conrelid = t.oid AND conindid IS NOT NULL
          )
        ORDER BY i.relname
    """
    with conn.cursor() as cur:
        cur.execute(query, (schema, table))
        rows = cur.fetchall()

    return [IndexDef(name=r[0], definition=r[1]) for r in rows]


def format_column(col: ColumnDef) -> str:
    """Format a single column definition."""
    parts = [f"    {quote_ident(col.name)}", col.data_type]

    # NOT NULL (but not for generated columns - it's implicit)
    if col.not_null and not col.generated:
        parts.append("NOT NULL")

    # Identity columns (PostgreSQL 10+)
    if col.identity == "a":
        parts.append("GENERATED ALWAYS AS IDENTITY")
    elif col.identity == "d":
        parts.append("GENERATED BY DEFAULT AS IDENTITY")
    # Generated columns (PostgreSQL 12+)
    elif col.generated == "s":
        # Default already contains the generation expression
        parts.append(f"GENERATED ALWAYS AS ({col.default}) STORED")
    elif col.default:
        parts.append(f"DEFAULT {col.default}")

    return " ".join(parts)


def quote_ident(name: str) -> str:
    """Quote identifier if needed (reserved words, special chars)."""
    # Simple heuristic: quote if not lowercase alphanumeric + underscore
    if name.isidentifier() and name.islower() and not is_reserved(name):
        return name
    return f'"{name}"'


def is_reserved(name: str) -> bool:
    """Check if name is a PostgreSQL reserved word (common ones)."""
    reserved = {
        "all",
        "analyse",
        "analyze",
        "and",
        "any",
        "array",
        "as",
        "asc",
        "asymmetric",
        "both",
        "case",
        "cast",
        "check",
        "collate",
        "column",
        "constraint",
        "create",
        "current_date",
        "current_role",
        "current_time",
        "current_timestamp",
        "current_user",
        "default",
        "deferrable",
        "desc",
        "distinct",
        "do",
        "else",
        "end",
        "except",
        "false",
        "fetch",
        "for",
        "foreign",
        "from",
        "grant",
        "group",
        "having",
        "in",
        "initially",
        "intersect",
        "into",
        "lateral",
        "leading",
        "limit",
        "localtime",
        "localtimestamp",
        "not",
        "null",
        "offset",
        "on",
        "only",
        "or",
        "order",
        "placing",
        "primary",
        "references",
        "returning",
        "select",
        "session_user",
        "some",
        "symmetric",
        "table",
        "then",
        "to",
        "trailing",
        "true",
        "union",
        "unique",
        "user",
        "using",
        "variadic",
        "when",
        "where",
        "window",
        "with",
    }
    return name.lower() in reserved


def format_table_name(schema: str, table: str) -> str:
    """Format table name, omitting 'public.' prefix."""
    if schema == "public":
        return quote_ident(table)
    return f"{quote_ident(schema)}.{quote_ident(table)}"


def strip_public_schema(definition: str) -> str:
    """Remove 'public.' prefix from DDL statements."""
    return definition.replace("public.", "").replace('"public".', "")


def generate_ddl(conn, schema: str, table: str) -> str:
    """Generate CREATE TABLE DDL with indexes."""
    if not table_exists(conn, schema, table):
        raise SchemaError(f"Table '{schema}.{table}' not found")

    columns = get_columns(conn, schema, table)
    constraints = get_constraints(conn, schema, table)
    indexes = get_indexes(conn, schema, table)

    if not columns:
        raise SchemaError(f"No columns found for '{schema}.{table}'")

    # Build CREATE TABLE
    table_name = format_table_name(schema, table)
    lines = [f"CREATE TABLE {table_name} ("]

    # Column definitions
    col_lines = [format_column(c) for c in columns]

    # Constraint definitions (strip public. from FK references)
    con_lines = [
        f"    CONSTRAINT {quote_ident(c.name)} {strip_public_schema(c.definition)}"
        for c in constraints
    ]

    # Combine columns and constraints with commas
    all_lines = col_lines + con_lines
    for i, line in enumerate(all_lines):
        if i < len(all_lines) - 1:
            lines.append(line + ",")
        else:
            lines.append(line)

    lines.append(");")

    # Add indexes as separate statements (strip public. from definitions)
    for idx in indexes:
        lines.append("")
        lines.append(strip_public_schema(idx.definition) + ";")

    return "\n".join(lines)


def run_schema(
    conn,
    table: str,
    schema: str = "public",
    json_output: bool = False,
) -> None:
    """Generate and print DDL."""
    ddl = generate_ddl(conn, schema, table)

    if json_output:
        print(json.dumps({"ddl": ddl}))
    else:
        print(ddl)
