from pocket_analyst._session import _Session
from pocket_analyst._db import connection


def connect_database(connection_string: str, session: _Session) -> str:
    result = connection.connect(connection_string, session)
    if result.startswith("Connected"):
        result += (
            "\n\nWould you like to add any context about this database — e.g. what it "
            "represents, key definitions, or known data quirks? "
            "Call set_context() to store it alongside the schema. This is optional."
        )
    return result


def disconnect_database(session: _Session) -> str:
    return connection.disconnect(session)


def set_database_context(context: str, session: _Session) -> str:
    connection.set_context(context, session)
    return f"Context saved.\n\n---\n{context}\n---"
