import unittest

from src.aiel_sdk.integrations.postgres.remote import RemotePostgres
from src.aiel_sdk.integrations.postgres.types import PostgresInsertRequest, PostgresQueryRequest


class _FakeIntegrationsService:
    def __init__(self):
        self.calls = []

    def invoke_action(self, workspace_id, project_id, connection_id, action, payload=None, request_id=None):
        self.calls.append(
            {
                "workspace_id": workspace_id,
                "project_id": project_id,
                "connection_id": connection_id,
                "action": action,
                "payload": payload,
                "request_id": request_id,
            }
        )
        if action == "postgres.query":
            return {"rows": [{"id": 1}], "row_count": 1}
        if action == "postgres.insert":
            return {"inserted": 1, "returning_rows": [{"id": 1}]}
        return {}


class _FakeClient:
    def __init__(self):
        self.integrations = _FakeIntegrationsService()


class RemotePostgresTests(unittest.TestCase):
    def test_query_invokes_expected_action(self):
        client = _FakeClient()
        pg = RemotePostgres(client=client, workspace_id="w1", project_id="p1", connection_id="c1")

        out = pg.query(PostgresQueryRequest(sql="select 1"), request_id="req_1")

        self.assertEqual(out.row_count, 1)
        self.assertEqual(len(out.rows), 1)
        call = client.integrations.calls[-1]
        self.assertEqual(call["action"], "postgres.query")
        self.assertEqual(call["payload"], {"params": {"sql": "select 1"}})
        self.assertEqual(call["request_id"], "req_1")

    def test_insert_invokes_expected_action(self):
        client = _FakeClient()
        pg = RemotePostgres(client=client, workspace_id="w1", project_id="p1", connection_id="c1")

        out = pg.insert(
            PostgresInsertRequest(table="users", values={"name": "Alden"}),
            request_id="req_2",
        )

        self.assertEqual(out.inserted, 1)
        self.assertEqual(len(out.returning_rows), 1)
        call = client.integrations.calls[-1]
        self.assertEqual(call["action"], "postgres.insert")
        self.assertEqual(call["payload"], {"params": {"table": "users", "values": {"name": "Alden"}}})
        self.assertEqual(call["request_id"], "req_2")

    def test_validation(self):
        with self.assertRaises(ValueError):
            PostgresQueryRequest(sql="").to_params()
        with self.assertRaises(ValueError):
            PostgresInsertRequest(table="", values={"x": 1}).to_params()
        with self.assertRaises(ValueError):
            PostgresInsertRequest(table="t", values={}).to_params()


if __name__ == "__main__":
    unittest.main(verbosity=2)
