"""
cassachange.baseline — Introspect a live Cassandra keyspace and emit a
versioned V{version}__baseline_{keyspace}.cql file.

Usage:
    cassachange baseline --keyspace myapp
    cassachange baseline --keyspace myapp --output ./migrations/ --baseline-version 1.0.0

What is captured (IF NOT EXISTS throughout for safe re-runs):
  • User-Defined Types (UDTs)
  • Tables — columns, partition/clustering keys, clustering order, comments
  • Secondary indexes and SAI (Storage Attached) indexes
  • Materialized Views — commented out with a note (not supported on AstraDB)
  • User-Defined Functions (UDFs)
  • User-Defined Aggregates (UDAs)

What is NOT captured (infrastructure — Terraform's responsibility):
  • Keyspace creation / replication strategy
  • Auth, roles, permissions
"""

import logging
import os
from datetime import datetime

log = logging.getLogger("cassachange.baseline")

SYSTEM_KEYSPACES = {
    "system", "system_auth", "system_distributed", "system_schema",
    "system_traces", "system_virtual_schema", "system_views",
    "data_endpoint_auth", "datastax_sla",
}


def generate_baseline(session, keyspace: str, output_dir: str,
                      version: str = "0.0.0") -> str:
    """
    Introspect *keyspace* and write the baseline script to *output_dir*.
    Returns the absolute path of the written file.
    """
    if keyspace.lower() in SYSTEM_KEYSPACES:
        raise ValueError(f"Cannot baseline system keyspace: {keyspace}")

    log.info("Introspecting keyspace '%s'...", keyspace)

    lines: list[str] = []
    lines += _header(keyspace, version)
    lines += _udts(session, keyspace)
    lines += _tables(session, keyspace)
    lines += _indexes(session, keyspace)
    lines += _views(session, keyspace)
    lines += _functions(session, keyspace)
    lines += _aggregates(session, keyspace)

    os.makedirs(output_dir, exist_ok=True)
    fname = f"V{version}__baseline_{keyspace}.cql"
    fpath = os.path.join(output_dir, fname)
    with open(fpath, "w", encoding="utf-8") as fh:
        fh.write("\n".join(lines))

    log.info("Baseline written → %s  (%d lines)", fpath, len(lines))
    return fpath


# ---------------------------------------------------------------------------
# Header
# ---------------------------------------------------------------------------

def _header(keyspace: str, version: str) -> list:
    ts = datetime.utcnow().strftime("%Y-%m-%d %H:%M:%S UTC")
    return [
        f"-- cassachange baseline: {keyspace}",
        f"-- Generated : {ts}",
        f"-- Version   : V{version}",
        f"--",
        f"-- Auto-generated by 'cassachange baseline'.",
        f"-- All statements use IF NOT EXISTS for safe re-execution.",
        f"-- Do NOT edit. Create incremental V__ scripts to evolve the schema.",
        "", "",
    ]


# ---------------------------------------------------------------------------
# User-Defined Types
# ---------------------------------------------------------------------------

def _udts(session, keyspace: str) -> list:
    rows = list(session.execute(
        "SELECT type_name, field_names, field_types "
        "FROM system_schema.types WHERE keyspace_name = %s", (keyspace,),
    ))
    if not rows:
        return []

    out = ["-- UDTs", "-- " + "-" * 70, ""]
    for r in sorted(rows, key=lambda x: x.type_name):
        fields = ",\n".join(
            f"    {_q(n)} {t}"
            for n, t in zip(r.field_names, r.field_types)
        )
        out += [
            f"CREATE TYPE IF NOT EXISTS {_qks(keyspace)}.{_q(r.type_name)} (",
            fields,
            ");", "",
        ]
    return out


# ---------------------------------------------------------------------------
# Tables
# ---------------------------------------------------------------------------

def _tables(session, keyspace: str) -> list:
    tables = list(session.execute(
        "SELECT table_name, comment FROM system_schema.tables "
        "WHERE keyspace_name = %s", (keyspace,),
    ))
    if not tables:
        return []

    out = ["-- Tables", "-- " + "-" * 70, ""]
    for tr in sorted(tables, key=lambda x: x.table_name):
        tname = tr.table_name
        cols  = list(session.execute(
            "SELECT column_name, kind, position, type, clustering_order "
            "FROM system_schema.columns "
            "WHERE keyspace_name = %s AND table_name = %s", (keyspace, tname),
        ))
        pk  = sorted([c for c in cols if c.kind == "partition_key"],  key=lambda c: c.position)
        ck  = sorted([c for c in cols if c.kind == "clustering"],      key=lambda c: c.position)
        reg = [c for c in cols if c.kind in ("regular", "static")]

        defs = []
        for c in pk + ck + reg:
            static = " STATIC" if c.kind == "static" else ""
            defs.append(f"    {_q(c.column_name)} {c.type}{static}")

        pk_str = ", ".join(_q(c.column_name) for c in pk)
        ck_str = ", ".join(_q(c.column_name) for c in ck)
        pk_clause = f"({pk_str}), {ck_str}" if ck else pk_str
        defs.append(f"    PRIMARY KEY ({pk_clause})")

        body = ",\n".join(defs)
        out.append(f"CREATE TABLE IF NOT EXISTS {_qks(keyspace)}.{_q(tname)} (")
        out.append(body)
        out.append(")")

        opts = []
        if ck:
            order = ", ".join(
                f"{_q(c.column_name)} {(c.clustering_order or 'ASC').upper()}"
                for c in ck
            )
            opts.append(f"CLUSTERING ORDER BY ({order})")
        if getattr(tr, "comment", None):
            opts.append(f"comment = '{tr.comment.replace(chr(39), chr(39)*2)}'")

        if opts:
            out[-1] = out[-1]  # keep the closing paren line
            out.append("WITH " + "\nAND ".join(opts))

        out[-1] += ";"
        out.append("")
    return out


# ---------------------------------------------------------------------------
# Indexes
# ---------------------------------------------------------------------------

def _indexes(session, keyspace: str) -> list:
    rows = list(session.execute(
        "SELECT index_name, table_name, kind, options "
        "FROM system_schema.indexes WHERE keyspace_name = %s", (keyspace,),
    ))
    if not rows:
        return []

    out = ["-- Indexes", "-- " + "-" * 70, ""]
    for r in sorted(rows, key=lambda x: (x.table_name, x.index_name)):
        opts   = r.options or {}
        target = opts.get("target", "")
        kind   = (r.kind or "").upper()
        tbl    = f"{_qks(keyspace)}.{_q(r.table_name)}"

        if kind == "CUSTOM":
            cls = opts.get("class_name", "")
            using = "StorageAttachedIndex" if ("StorageAttachedIndex" in cls or "sai" in cls.lower()) else cls
            out.append(
                f"CREATE CUSTOM INDEX IF NOT EXISTS {_q(r.index_name)} "
                f"ON {tbl} ({target}) USING '{using}';"
            )
        else:
            out.append(
                f"CREATE INDEX IF NOT EXISTS {_q(r.index_name)} "
                f"ON {tbl} ({target});"
            )
        out.append("")
    return out


# ---------------------------------------------------------------------------
# Materialized Views (noted, not reconstructed — AstraDB doesn't support them)
# ---------------------------------------------------------------------------

def _views(session, keyspace: str) -> list:
    try:
        rows = list(session.execute(
            "SELECT view_name, base_table_name FROM system_schema.views "
            "WHERE keyspace_name = %s", (keyspace,),
        ))
    except Exception:
        return []
    if not rows:
        return []

    out = ["-- Materialized Views (NOT supported on AstraDB Serverless)",
           "-- " + "-" * 70,
           "-- These views were detected but DDL cannot be reliably reconstructed.",
           "-- Recreate them manually or replace with real tables + SAI indexes.", ""]
    for r in sorted(rows, key=lambda x: x.view_name):
        out.append(f"-- View: {r.view_name}  (base: {r.base_table_name})")
    out.append("")
    return out


# ---------------------------------------------------------------------------
# UDFs / UDAs
# ---------------------------------------------------------------------------

def _functions(session, keyspace: str) -> list:
    try:
        rows = list(session.execute(
            "SELECT function_name, argument_types, argument_names, return_type, "
            "called_on_null_input, language, body "
            "FROM system_schema.functions WHERE keyspace_name = %s", (keyspace,),
        ))
    except Exception:
        return []
    if not rows:
        return []

    out = ["-- User-Defined Functions", "-- " + "-" * 70, ""]
    for r in sorted(rows, key=lambda x: x.function_name):
        args = ", ".join(
            f"{_q(n)} {t}" for n, t in zip(r.argument_names or [], r.argument_types or [])
        )
        null_behavior = "CALLED ON NULL INPUT" if r.called_on_null_input else "RETURNS NULL ON NULL INPUT"
        out += [
            f"CREATE OR REPLACE FUNCTION {_qks(keyspace)}.{_q(r.function_name)} ({args})",
            f"    {null_behavior}",
            f"    RETURNS {r.return_type}",
            f"    LANGUAGE {r.language}",
            f"    AS $${r.body}$$;",
            "",
        ]
    return out


def _aggregates(session, keyspace: str) -> list:
    try:
        rows = list(session.execute(
            "SELECT aggregate_name, argument_types, state_func, state_type, "
            "final_func, return_type, initcond "
            "FROM system_schema.aggregates WHERE keyspace_name = %s", (keyspace,),
        ))
    except Exception:
        return []
    if not rows:
        return []

    out = ["-- User-Defined Aggregates", "-- " + "-" * 70, ""]
    for r in sorted(rows, key=lambda x: x.aggregate_name):
        arg_types = ", ".join(r.argument_types or [])
        lines = [
            f"CREATE OR REPLACE AGGREGATE {_qks(keyspace)}.{_q(r.aggregate_name)} ({arg_types})",
            f"    SFUNC {r.state_func}",
            f"    STYPE {r.state_type}",
        ]
        if r.final_func:
            lines.append(f"    FINALFUNC {r.final_func}")
        if r.initcond:
            lines.append(f"    INITCOND {r.initcond}")
        lines.append(";")
        out += lines
        out.append("")
    return out


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

_RESERVED = {
    "table","index","type","function","aggregate","view","role","user","key",
    "token","order","limit","group","where","select","insert","update","delete",
    "create","drop","alter","with","from","into","set","using","if","not",
    "exists","and","or","in","is","null","true","false","primary","partition",
    "clustering","static","list","map",
}

def _q(name: str) -> str:
    if name.lower() in _RESERVED or not name.replace("_","").isalnum():
        return f'"{name}"'
    return name

def _qks(ks: str) -> str:
    return _q(ks)
