from sqlalchemy import create_engine, text
from sqlalchemy.engine import Engine

from pocket_analyst._session import _Session


def connect(connection_string: str, session: _Session) -> str:
    try:
        if session.engine is not None:
            session.engine.dispose()
        if connection_string.startswith("duckdb://"):
            engine = create_engine(connection_string, connect_args={"read_only": True})
        else:
            engine = create_engine(connection_string)
        with engine.connect() as conn:
            conn.execute(text("SELECT 1"))
        session.engine = engine
        session.connection_string = connection_string
        session.user_context = None
        return f"Connected successfully ({engine.dialect.name})"
    except Exception as e:
        return f"Connection failed: {e}"


def disconnect(session: _Session) -> str:
    if session.engine is None:
        return "No database was connected."
    session.engine.dispose()
    session.engine = None
    session.connection_string = None
    session.user_context = None
    return "Disconnected. All database state and context have been cleared."


def set_context(text_val: str, session: _Session) -> None:
    session.user_context = text_val


def get_context(session: _Session) -> str | None:
    return session.user_context


def get_engine(session: _Session) -> Engine:
    if session.engine is None:
        raise RuntimeError("No database connected. Call connect() first.")
    return session.engine


def is_connected(session: _Session) -> bool:
    return session.engine is not None
