#!/usr/bin/env python3

from __future__ import annotations

import os
import ssl
import sys
import time
import json
from dataclasses import dataclass
from datetime import timedelta

from cassandra import ConsistencyLevel
from cassandra.auth import PlainTextAuthProvider
from cassandra.cluster import Cluster
from cassandra.query import SimpleStatement


DEFAULT_HOSTS = "127.0.0.1"
DEFAULT_PORT = 9042
DEFAULT_USERNAME = "cassandra"
DEFAULT_PASSWORD = "cassandra"
DEFAULT_TIMEOUT = timedelta(seconds=5)
DEFAULT_CONNECT_TIMEOUT = timedelta(seconds=5)
DEFAULT_PROTO_VERSION = 4
DEFAULT_KEYSPACE = "sample_scylla_driver_connectivity"
DEFAULT_TABLE = "items"
DEFAULT_ENV_FILE = ".env"


@dataclass(frozen=True)
class SampleConfig:
    hosts: list[str]
    port: int
    username: str
    password: str
    timeout_seconds: float
    connect_timeout_seconds: float
    protocol_version: int
    tls_enabled: bool
    tls_skip_verify: bool
    tls_server_name: str | None
    keyspace: str
    table: str


def main() -> int:
    started_at = time.monotonic()
    try:
        load_env_file(env("CASSANDRA_ENV_FILE", DEFAULT_ENV_FILE))
        config = config_from_env()
        print_connection_input(config)
        session = connect(config)
    except Exception as exc:
        print(f"connect with ScyllaDB Python driver failed: {exc}", file=sys.stderr)
        return 1

    try:
        local = query_system_local(session, config.timeout_seconds)
        results = run_smoke_tests(session, config)
    except Exception as exc:
        print(f"ScyllaDB Python driver smoke test failed: {exc}", file=sys.stderr)
        return 1
    finally:
        session.cluster.shutdown()

    elapsed_ms = round((time.monotonic() - started_at) * 1000)
    print("ScyllaDB Python driver connectivity OK")
    print(f"hosts: {','.join(config.hosts)}")
    print(f"port: {config.port}")
    print(f"protocol_version: {config.protocol_version}")
    print(f"keyspace: {config.keyspace}")
    print(f"table: {config.table}")
    print(f"elapsed: {elapsed_ms}ms")
    print(f"cluster_name: {local.cluster_name}")
    print(f"data_center: {local.data_center}")
    print(f"host_id: {local.host_id}")
    print(f"release_version: {local.release_version}")
    print(f"rpc_address: {local.rpc_address}")
    print(f"rpc_port: {local.rpc_port}")
    for result in results:
        print(f"summary: {result}")
    return 0


def connect(config: SampleConfig):
    auth_provider = PlainTextAuthProvider(
        username=config.username,
        password=config.password,
    )
    ssl_context = build_ssl_context(config) if config.tls_enabled else None
    ssl_options = None
    if config.tls_enabled and config.tls_server_name:
        ssl_options = {"server_hostname": config.tls_server_name}
    cluster = Cluster(
        contact_points=config.hosts,
        port=config.port,
        auth_provider=auth_provider,
        protocol_version=config.protocol_version,
        connect_timeout=config.connect_timeout_seconds,
        control_connection_timeout=config.connect_timeout_seconds,
        ssl_context=ssl_context,
        ssl_options=ssl_options,
    )
    try:
        return cluster.connect(wait_for_all_pools=True)
    except Exception:
        cluster.shutdown()
        raise


def print_connection_input(config: SampleConfig) -> None:
    print("connection input:", flush=True)
    print(
        format_value(
            {
                "hosts": config.hosts,
                "port": config.port,
                "username": config.username,
                "password": "<redacted>" if config.password else "",
                "protocol_version": config.protocol_version,
                "timeout_seconds": config.timeout_seconds,
                "connect_timeout_seconds": config.connect_timeout_seconds,
                "tls_enabled": config.tls_enabled,
                "tls_skip_verify": config.tls_skip_verify,
                "tls_server_name": config.tls_server_name,
                "keyspace": config.keyspace,
                "table": config.table,
            }
        ),
        flush=True,
    )


def query_system_local(session, timeout_seconds: float):
    statement = session.prepare(
        """
        SELECT cluster_name, data_center, host_id, release_version, rpc_address, rpc_port
        FROM system.local
        WHERE key = ?
        """
    )
    statement.consistency_level = ConsistencyLevel.ONE
    result = session.execute(statement, ("local",), timeout=timeout_seconds)
    row = result.one()
    if row is None:
        raise RuntimeError("system.local returned no rows")
    return row


def run_smoke_tests(session, config: SampleConfig) -> list[str]:
    keyspace = quote_identifier(config.keyspace)
    table = quote_identifier(config.table)
    full_table = f"{keyspace}.{table}"
    timeout = config.timeout_seconds
    results: list[str] = []

    execute(session, f"DROP KEYSPACE IF EXISTS {keyspace}", timeout)

    case_input = normalize_cql(
        f"""
        CREATE KEYSPACE {keyspace}
        WITH replication = {{ 'class': 'SimpleStrategy', 'replication_factor': 1 }}
        """
    )
    print_case_input(1, "create keyspace", case_input)
    execute(session, case_input, timeout)
    print_case_result({"status": "ok", "keyspace": config.keyspace})
    results.append("create keyspace")

    case_input = normalize_cql(
        f"""
        CREATE TABLE {full_table} (
            id int PRIMARY KEY,
            name text,
            score int
        )
        """
    )
    print_case_input(2, "create table", case_input)
    execute(session, case_input, timeout)
    print_case_result({"status": "ok", "table": full_table})
    results.append("create table")

    case_input = f"DESC KEYSPACE {keyspace}"
    print_case_input(3, "desc keyspace", case_input)
    keyspace_desc = list(execute(session, case_input, timeout))
    if len(keyspace_desc) == 0:
        raise RuntimeError("DESC KEYSPACE returned no rows")
    print_case_result({"row_count": len(keyspace_desc), "rows": rows_to_dicts(keyspace_desc)})
    results.append(f"desc keyspace rows={len(keyspace_desc)}")

    case_input = f"DESC TABLE {full_table}"
    print_case_input(4, "desc keyspace.table", case_input)
    table_desc = list(execute(session, case_input, timeout))
    if len(table_desc) == 0:
        raise RuntimeError("DESC keyspace.table returned no rows")
    print_case_result({"row_count": len(table_desc), "rows": rows_to_dicts(table_desc)})
    results.append(f"desc keyspace.table rows={len(table_desc)}")

    insert_cql = f"INSERT INTO {full_table} (id, name, score) VALUES (?, ?, ?)"
    insert_rows = [
        {"id": item_id, "name": f"name-{item_id}", "score": item_id * 10}
        for item_id in range(1, 6)
    ]
    print_case_input(5, "insert 5 rows", {"cql": insert_cql, "rows": insert_rows})
    insert_stmt = session.prepare(insert_cql)
    insert_stmt.consistency_level = ConsistencyLevel.ONE
    for row in insert_rows:
        session.execute(
            insert_stmt,
            (row["id"], row["name"], row["score"]),
            timeout=timeout,
        )
    print_case_result({"status": "ok", "inserted_rows": insert_rows})
    results.append("insert 5 rows")

    select_by_pk_cql = f"SELECT id, name, score FROM {full_table} WHERE id = ?"
    print_case_input(6, "select by pk", {"cql": select_by_pk_cql, "params": [3]})
    select_by_pk = session.prepare(select_by_pk_cql)
    select_by_pk.consistency_level = ConsistencyLevel.ONE
    pk_row = session.execute(select_by_pk, (3,), timeout=timeout).one()
    if pk_row is None:
        raise RuntimeError("SELECT by pk returned no row")
    if (pk_row.id, pk_row.name, pk_row.score) != (3, "name-3", 30):
        raise RuntimeError(f"SELECT by pk returned unexpected row: {pk_row}")
    print_case_result({"row": row_to_dict(pk_row)})
    results.append("select by pk")

    delete_cql = f"DELETE FROM {full_table} WHERE id = ?"
    select_limit_cql = f"SELECT id, name, score FROM {full_table} LIMIT 5"
    print_case_input(
        7,
        "delete 2 rows and select limit 5",
        {
            "delete": {"cql": delete_cql, "params": [[1], [2]]},
            "select": select_limit_cql,
        },
    )
    delete_stmt = session.prepare(delete_cql)
    delete_stmt.consistency_level = ConsistencyLevel.ONE
    for item_id in (1, 2):
        session.execute(delete_stmt, (item_id,), timeout=timeout)
    rows = list(execute(session, select_limit_cql, timeout))
    ids = sorted(row.id for row in rows)
    if ids != [3, 4, 5]:
        raise RuntimeError(f"SELECT * LIMIT 5 after delete returned ids {ids}, want [3, 4, 5]")
    print_case_result(
        {
            "deleted_ids": [1, 2],
            "selected_row_count": len(rows),
            "selected_rows": rows_to_dicts(rows),
        }
    )
    results.append(f"delete 2 rows and select limit 5 rows={len(rows)}")

    case_input = f"DROP TABLE {full_table}"
    print_case_input(8, "drop table", case_input)
    execute(session, case_input, timeout)
    print_case_result({"status": "ok", "table": full_table})
    results.append("drop table")

    case_input = f"DROP KEYSPACE {keyspace}"
    print_case_input(9, "drop keyspace", case_input)
    execute(session, case_input, timeout)
    print_case_result({"status": "ok", "keyspace": config.keyspace})
    results.append("drop keyspace")

    return results


def execute(session, statement: str, timeout_seconds: float):
    simple = SimpleStatement(statement, consistency_level=ConsistencyLevel.ONE)
    return session.execute(simple, timeout=timeout_seconds)


def print_case_input(index: int, name: str, value) -> None:
    print(f"\ncase {index}: {name}", flush=True)
    print("input:", flush=True)
    print(format_value(value), flush=True)


def print_case_result(value) -> None:
    print("result:", flush=True)
    print(format_value(value), flush=True)


def format_value(value) -> str:
    if isinstance(value, str):
        return value
    return json.dumps(to_jsonable(value), indent=2, sort_keys=True)


def rows_to_dicts(rows) -> list[dict]:
    return [row_to_dict(row) for row in rows]


def row_to_dict(row) -> dict:
    if hasattr(row, "_asdict"):
        return {key: to_jsonable(value) for key, value in row._asdict().items()}
    return {"value": to_jsonable(row)}


def to_jsonable(value):
    if isinstance(value, dict):
        return {str(key): to_jsonable(item) for key, item in value.items()}
    if isinstance(value, (list, tuple, set)):
        return [to_jsonable(item) for item in value]
    if isinstance(value, (str, int, float, bool)) or value is None:
        return value
    return str(value)


def normalize_cql(statement: str) -> str:
    return "\n".join(line.strip() for line in statement.strip().splitlines())


def config_from_env() -> SampleConfig:
    hosts = split_hosts(env("CASSANDRA_HOSTS", DEFAULT_HOSTS))
    if not hosts:
        raise ValueError("CASSANDRA_HOSTS must include at least one host")

    return SampleConfig(
        hosts=hosts,
        port=int_env("CASSANDRA_PORT", DEFAULT_PORT),
        username=env("CASSANDRA_USERNAME", DEFAULT_USERNAME),
        password=env("CASSANDRA_PASSWORD", DEFAULT_PASSWORD),
        timeout_seconds=duration_env("CASSANDRA_TIMEOUT", DEFAULT_TIMEOUT),
        connect_timeout_seconds=duration_env(
            "CASSANDRA_CONNECT_TIMEOUT",
            DEFAULT_CONNECT_TIMEOUT,
        ),
        protocol_version=int_env("CASSANDRA_PROTO_VERSION", DEFAULT_PROTO_VERSION),
        tls_enabled=bool_env("CASSANDRA_TLS", False),
        tls_skip_verify=bool_env("CASSANDRA_TLS_SKIP_VERIFY", False),
        tls_server_name=os.getenv("CASSANDRA_TLS_SERVER_NAME") or None,
        keyspace=identifier_env("CASSANDRA_SAMPLE_KEYSPACE", DEFAULT_KEYSPACE),
        table=identifier_env("CASSANDRA_SAMPLE_TABLE", DEFAULT_TABLE),
    )


def build_ssl_context(config: SampleConfig) -> ssl.SSLContext:
    context = ssl.create_default_context()
    if config.tls_skip_verify:
        context.check_hostname = False
        context.verify_mode = ssl.CERT_NONE
    return context


def split_hosts(value: str) -> list[str]:
    hosts: list[str] = []
    for field in value.replace(";", ",").replace(" ", ",").split(","):
        field = field.strip()
        if field:
            hosts.append(field)
    return hosts


def env(key: str, fallback: str) -> str:
    value = os.getenv(key, "").strip()
    return value or fallback


def load_env_file(path: str) -> None:
    if not path or not os.path.exists(path):
        return
    with open(path, encoding="utf-8") as file:
        for raw_line in file:
            line = raw_line.strip()
            if not line or line.startswith("#"):
                continue
            if line.startswith("export "):
                line = line[len("export ") :].strip()
            if "=" not in line:
                continue
            key, value = line.split("=", 1)
            key = key.strip()
            value = clean_env_value(value)
            if key:
                os.environ[key] = value


def clean_env_value(value: str) -> str:
    value = value.strip()
    if value.endswith("\\"):
        value = value[:-1].strip()
    if len(value) >= 2 and value[0] == value[-1] and value[0] in {"'", '"'}:
        value = value[1:-1]
    return value


def int_env(key: str, fallback: int) -> int:
    value = os.getenv(key, "").strip()
    if not value:
        return fallback
    try:
        return int(value)
    except ValueError as exc:
        raise ValueError(f"{key} must be an integer") from exc


def identifier_env(key: str, fallback: str) -> str:
    value = env(key, fallback)
    validate_identifier(key, value)
    return value


def validate_identifier(key: str, value: str) -> None:
    if not value:
        raise ValueError(f"{key} must not be empty")
    if not (value[0].isalpha() or value[0] == "_"):
        raise ValueError(f"{key} must start with a letter or underscore")
    if any(not (ch.isalnum() or ch == "_") for ch in value):
        raise ValueError(f"{key} must contain only letters, digits, or underscores")


def quote_identifier(value: str) -> str:
    validate_identifier("identifier", value)
    return value


def bool_env(key: str, fallback: bool) -> bool:
    value = os.getenv(key, "").strip().lower()
    if not value:
        return fallback
    if value in {"1", "true", "t", "yes", "y", "on"}:
        return True
    if value in {"0", "false", "f", "no", "n", "off"}:
        return False
    raise ValueError(f"{key} must be a boolean")


def duration_env(key: str, fallback: timedelta) -> float:
    value = os.getenv(key, "").strip()
    if not value:
        return fallback.total_seconds()
    return parse_duration(value)


def parse_duration(value: str) -> float:
    units = {
        "ms": 0.001,
        "s": 1,
        "m": 60,
    }
    for suffix, multiplier in units.items():
        if value.endswith(suffix):
            number = value[: -len(suffix)]
            try:
                return float(number) * multiplier
            except ValueError as exc:
                raise ValueError(f"duration {value!r} has invalid numeric value") from exc
    try:
        return float(value)
    except ValueError as exc:
        raise ValueError(
            f"duration {value!r} must look like 5s, 500ms, 1m, or a seconds value"
        ) from exc


if __name__ == "__main__":
    raise SystemExit(main())