import datetime
import time
import uuid

import pytest
import sqlalchemy as sa
from sqlalchemy.exc import DBAPIError

# Public API
from dbos import DBOS, GetWorkflowsInput, SetWorkflowID


def test_transaction_errors(dbos: DBOS) -> None:
    retry_counter: int = 0

    @DBOS.transaction()
    def test_retry_transaction(max_retry: int) -> int:
        nonlocal retry_counter
        if retry_counter < max_retry:
            retry_counter += 1
            base_err = BaseException()
            base_err.pgcode = "40001"  # type: ignore
            err = DBAPIError("Serialization test error", {}, base_err)
            raise err
        return max_retry

    @DBOS.transaction()
    def test_noretry_transaction() -> None:
        nonlocal retry_counter
        retry_counter += 1
        DBOS.sql_session.execute(sa.text("selct abc from c;")).fetchall()

    res = test_retry_transaction(10)
    assert res == 10
    assert retry_counter == 10

    with pytest.raises(Exception) as exc_info:
        test_noretry_transaction()
    assert exc_info.value.orig.pgcode == "42601"  # type: ignore
    assert retry_counter == 11


def test_notification_errors(dbos: DBOS) -> None:
    @DBOS.workflow()
    def test_send_workflow(dest_uuid: str, topic: str) -> str:
        DBOS.send(dest_uuid, "test1")
        DBOS.send(dest_uuid, "test2", topic=topic)
        DBOS.send(dest_uuid, "test3")
        return dest_uuid

    @DBOS.workflow()
    def test_recv_workflow(topic: str) -> str:
        msg1 = DBOS.recv(topic, timeout_seconds=10)
        msg2 = DBOS.recv(timeout_seconds=10)
        msg3 = DBOS.recv(timeout_seconds=10)
        return "-".join([str(msg1), str(msg2), str(msg3)])

    # Crash the notification connection and make sure send/recv works on time.
    while dbos.sys_db.notification_conn is None:
        time.sleep(1)
    dbos.sys_db.notification_conn.close()
    assert dbos.sys_db.notification_conn.closed == 1

    # Wait for the connection to be re-established
    while dbos.sys_db.notification_conn.closed != 0:
        time.sleep(1)

    dest_uuid = str("sruuid1")
    with SetWorkflowID(dest_uuid):
        handle = dbos.start_workflow(test_recv_workflow, "testtopic")
        assert handle.get_workflow_id() == dest_uuid

    send_uuid = str("sruuid2")
    with SetWorkflowID(send_uuid):
        res = test_send_workflow(handle.get_workflow_id(), "testtopic")
        assert res == dest_uuid

    begin_time = time.time()
    assert handle.get_result() == "test2-test1-test3"
    duration = time.time() - begin_time
    assert duration < 3.0


def test_buffer_flush_errors(dbos: DBOS) -> None:
    @DBOS.transaction()
    def test_transaction(var: str) -> str:
        rows = DBOS.sql_session.execute(sa.text("SELECT 1")).fetchall()
        return var + str(rows[0][0])

    cur_time: str = datetime.datetime.now().isoformat()
    gwi: GetWorkflowsInput = GetWorkflowsInput()
    gwi.start_time = cur_time

    res = test_transaction("bob")
    assert res == "bob1"

    dbos.sys_db.wait_for_buffer_flush()
    wfs = dbos.sys_db.get_workflows(gwi)
    assert len(wfs.workflow_uuids) == 1

    # Crash the system database connection and make sure the buffer flush works on time.
    backup_engine = dbos.sys_db.engine
    dbos.sys_db.engine = sa.create_engine(
        "postgresql+psycopg2://fake:database@localhost/fake_db"
    )

    res = test_transaction("bob")
    assert res == "bob1"

    # Should see some errors in the logs
    time.sleep(2)

    # Switch back to the original good engine.
    dbos.sys_db.engine = backup_engine

    dbos.sys_db.wait_for_buffer_flush()
    wfs = dbos.sys_db.get_workflows(gwi)
    assert len(wfs.workflow_uuids) == 2
