"""Units for the Postgres and Stripe connectors, plus the SQL safety guard."""

from __future__ import annotations

import uuid

import httpx
import pytest
from legate.agent.tooling import ToolContext
from legate.config import get_settings
from legate.connectors.mysql import MySQLConnector
from legate.connectors.postgres import PostgresConnector
from legate.connectors.registry import catalog_tool_specs, get_connector, known_tool_names
from legate.connectors.sql_base import SqlSafetyError, _ensure_read_only, _single_statement
from legate.connectors.stripe import StripeConnector


def _ctx() -> ToolContext:
    return ToolContext(workspace_id=uuid.uuid4(), settings=get_settings())


# --- SQL safety guard ---------------------------------------------------------


@pytest.mark.parametrize(
    "sql",
    [
        "SELECT * FROM users",
        "  select id from t where x = %s ",
        "WITH recent AS (SELECT * FROM orders) SELECT * FROM recent",
        "select 1; -- trailing comment and semicolon ok",
    ],
)
def test_read_only_guard_allows_selects(sql: str) -> None:
    assert _ensure_read_only(sql)  # no raise


@pytest.mark.parametrize(
    "sql",
    [
        "DELETE FROM users",
        "UPDATE users SET admin = true",
        "SELECT 1; DROP TABLE users",  # statement chaining
        "WITH x AS (DELETE FROM t RETURNING *) SELECT * FROM x",  # data-modifying CTE
        "INSERT INTO t VALUES (1)",
        "select * from t /* sneaky */ ; truncate t",
    ],
)
def test_read_only_guard_blocks_writes(sql: str) -> None:
    with pytest.raises(SqlSafetyError):
        _ensure_read_only(sql)


def test_single_statement_rejects_chaining() -> None:
    assert _single_statement("update t set a = 1") == "update t set a = 1"
    with pytest.raises(SqlSafetyError):
        _single_statement("update t set a = 1; delete from t")
    with pytest.raises(SqlSafetyError):
        _single_statement("   ")


# --- Postgres connector -------------------------------------------------------


async def test_postgres_query_runs_and_caps_rows() -> None:
    conn = PostgresConnector()
    captured: dict = {}

    def fake_query(config, secrets, sql, params, max_rows):
        captured.update(sql=sql, params=params, max_rows=max_rows)
        return {"columns": ["id", "name"], "rows": [[1, "a"], [2, "b"]]}

    conn._run_query = fake_query
    res = await conn.execute(
        "postgres.query",
        {"sql": "SELECT id, name FROM t", "params": ["x"], "max_rows": 99999},
        {"host": "h", "database": "d", "user": "u"},
        {"password": "p"},
        _ctx(),
    )
    assert res.ok and res.data["row_count"] == 2 and res.data["columns"] == ["id", "name"]
    assert captured["max_rows"] == 1000  # capped to the hard maximum
    assert captured["params"] == ["x"]


async def test_postgres_query_rejects_write() -> None:
    conn = PostgresConnector()
    conn._run_query = lambda *a, **k: pytest.fail("should not reach the database")
    res = await conn.execute(
        "postgres.query",
        {"sql": "DELETE FROM t"},
        {"host": "h", "database": "d", "user": "u"},
        {"password": "p"},
        _ctx(),
    )
    assert not res.ok and "SELECT" in res.error


async def test_postgres_execute_writes_and_blocks_chaining() -> None:
    conn = PostgresConnector()
    conn._run_exec = lambda config, secrets, sql, params: 3
    ok = await conn.execute(
        "postgres.execute",
        {"sql": "UPDATE t SET a = %s", "params": [1]},
        {"host": "h", "database": "d", "user": "u"},
        {"password": "p"},
        _ctx(),
    )
    assert ok.ok and ok.data["rows_affected"] == 3

    chained = await conn.execute(
        "postgres.execute",
        {"sql": "UPDATE t SET a = 1; DROP TABLE t"},
        {"host": "h", "database": "d", "user": "u"},
        {"password": "p"},
        _ctx(),
    )
    assert not chained.ok and "single" in chained.error


async def test_postgres_test_connection_uses_seam() -> None:
    conn = PostgresConnector()
    conn._run_query = lambda *a, **k: {"columns": ["?column?"], "rows": [[1]]}
    ok = await conn.test_connection({"host": "h", "database": "d", "user": "u"}, {"password": "p"})
    assert ok.ok

    def boom(*a, **k):
        raise OSError("refused")

    conn._run_query = boom
    bad = await conn.test_connection({"host": "h", "database": "d", "user": "u"}, {"password": "p"})
    assert not bad.ok and "failed" in bad.message


def test_postgres_conninfo_includes_credentials() -> None:
    conn = PostgresConnector()
    info = conn._conninfo(
        {"host": "db", "port": 5433, "database": "app", "user": "svc", "sslmode": "require"},
        {"password": "pw"},
    )
    assert (
        "host=db" in info
        and "dbname=app" in info
        and "password=pw" in info
        and "sslmode=require" in info
    )


# --- Stripe connector ---------------------------------------------------------


def _stripe_handler(captured: dict):
    def handler(request: httpx.Request) -> httpx.Response:
        captured["auth"] = request.headers.get("authorization")
        captured["path"] = request.url.path
        captured["method"] = request.method
        if request.method == "POST":
            captured["body"] = request.content.decode()
        if request.url.path.endswith("/balance"):
            return httpx.Response(200, json={"object": "balance"})
        if "/customers/" in request.url.path:
            return httpx.Response(200, json={"id": "cus_1", "email": "a@b.com", "name": "Ada"})
        if request.url.path.endswith("/charges"):
            return httpx.Response(
                200,
                json={
                    "data": [
                        {
                            "id": "ch_1",
                            "amount": 500,
                            "currency": "usd",
                            "status": "succeeded",
                            "paid": True,
                        }
                    ]
                },
            )
        if request.url.path.endswith("/refunds"):
            return httpx.Response(200, json={"id": "re_1", "status": "succeeded", "amount": 500})
        return httpx.Response(404, json={})

    return handler


async def test_stripe_get_customer_and_charges(set_connector_http):
    captured: dict = {}
    set_connector_http(_stripe_handler(captured))
    conn = StripeConnector()

    cust = await conn.execute(
        "stripe.get_customer", {"customer_id": "cus_1"}, {}, {"api_key": "sk_test"}, _ctx()
    )
    assert cust.ok and cust.data["email"] == "a@b.com"
    assert captured["auth"] == "Bearer sk_test"

    charges = await conn.execute(
        "stripe.list_charges", {"customer": "cus_1"}, {}, {"api_key": "sk_test"}, _ctx()
    )
    assert charges.ok and charges.data["charges"][0]["id"] == "ch_1"


async def test_stripe_create_refund_form_encoded(set_connector_http):
    captured: dict = {}
    set_connector_http(_stripe_handler(captured))
    conn = StripeConnector()

    res = await conn.execute(
        "stripe.create_refund",
        {"charge": "ch_1", "amount": 500},
        {},
        {"api_key": "sk_test"},
        _ctx(),
    )
    assert res.ok and res.data["id"] == "re_1" and res.data["status"] == "succeeded"
    assert captured["method"] == "POST" and captured["path"].endswith("/refunds")
    assert "charge=ch_1" in captured["body"] and "amount=500" in captured["body"]


async def test_stripe_refund_requires_target_and_key(set_connector_http):
    set_connector_http(_stripe_handler({}))
    conn = StripeConnector()

    no_target = await conn.execute(
        "stripe.create_refund", {"amount": 100}, {}, {"api_key": "sk_test"}, _ctx()
    )
    assert not no_target.ok and "required" in no_target.error

    no_key = await conn.execute("stripe.get_customer", {"customer_id": "cus_1"}, {}, {}, _ctx())
    assert not no_key.ok and "api_key" in no_key.error


async def test_stripe_test_connection(set_connector_http):
    set_connector_http(_stripe_handler({}))
    conn = StripeConnector()
    check = await conn.test_connection({}, {"api_key": "sk_test"})
    assert check.ok


# --- MySQL connector (shares BaseSqlConnector) --------------------------------


async def test_mysql_query_and_execute_share_the_guard() -> None:
    conn = MySQLConnector()
    conn._run_query = lambda *a, **k: {"columns": ["n"], "rows": [[1]]}
    conn._run_exec = lambda config, secrets, sql, params: 5

    read = await conn.execute(
        "mysql.query",
        {"sql": "SELECT n FROM t"},
        {"host": "h", "database": "d", "user": "u"},
        {"password": "p"},
        _ctx(),
    )
    assert read.ok and read.data["row_count"] == 1

    write = await conn.execute(
        "mysql.execute",
        {"sql": "DELETE FROM t WHERE id = %s", "params": [1]},
        {"host": "h", "database": "d", "user": "u"},
        {"password": "p"},
        _ctx(),
    )
    assert write.ok and write.data["rows_affected"] == 5

    blocked = await conn.execute(
        "mysql.query",
        {"sql": "UPDATE t SET a = 1"},
        {"host": "h", "database": "d", "user": "u"},
        {"password": "p"},
        _ctx(),
    )
    assert not blocked.ok and "SELECT" in blocked.error


def test_mysql_tool_names_and_sensitivity() -> None:
    specs = {s.name: s for s in MySQLConnector().list_tools()}
    assert set(specs) == {"mysql.query", "mysql.execute"}
    assert specs["mysql.query"].is_sensitive is False
    assert specs["mysql.execute"].is_sensitive is True


# --- Registry -----------------------------------------------------------------


def test_data_connectors_registered_with_correct_sensitivity():
    for t in ("postgres", "stripe", "mysql", "google"):
        assert get_connector(t) is not None

    catalog = {s.name: s for s in catalog_tool_specs()}
    assert {
        "postgres.query",
        "postgres.execute",
        "stripe.create_refund",
        "mysql.query",
        "gsheets.read_range",
        "gdrive.list_files",
    } <= set(catalog)
    assert {"postgres.query", "stripe.get_customer", "mysql.execute"} <= known_tool_names()

    assert catalog["postgres.query"].is_sensitive is False
    assert catalog["postgres.execute"].is_sensitive is True
    assert catalog["stripe.get_customer"].is_sensitive is False
    assert catalog["stripe.create_refund"].is_sensitive is True
    assert catalog["gsheets.read_range"].is_sensitive is False
    assert catalog["gsheets.append_row"].is_sensitive is True
