"""
pocket-analyst — embeddable database exploration toolkit.

Quick start:
    from pocket_analyst import PocketAnalyst

    pa = PocketAnalyst()
    pa.connect("duckdb:////path/to/mydb.duckdb")
    print(pa.explore_schema())
    result = pa.run_query("SELECT * FROM customers LIMIT 10")
    pa.disconnect()

Context manager:
    with PocketAnalyst() as pa:
        pa.connect("postgresql://user:pass@host/dbname")
        print(pa.list_tables())
"""

from pocket_analyst._session import _Session
from pocket_analyst._tools import connect as _connect
from pocket_analyst._tools import explore as _explore
from pocket_analyst._tools import query as _query


class PocketAnalyst:
    """Embedded database exploration client.

    Manages a single database session. Create one instance per connection.
    """

    def __init__(self):
        self._session = _Session()

    # ── Connection ──────────────────────────────────────────────────────────

    def connect(self, connection_string: str) -> str:
        """Connect to a database.

        Supported formats:
            duckdb:////path/to/mydb.duckdb
            sqlite:///path/to/mydb.sqlite
            postgresql://user:pass@host/dbname
            mysql+pymysql://user:pass@host/dbname
        """
        return _connect.connect_database(connection_string, self._session)

    def disconnect(self) -> str:
        """Close the current connection and clear session state."""
        return _connect.disconnect_database(self._session)

    def set_context(self, context: str) -> str:
        """Store free-text notes about this database.

        The context is prepended to schema output — useful for giving
        an LLM agent background on what the data represents.
        """
        return _connect.set_database_context(context, self._session)

    # ── Schema ──────────────────────────────────────────────────────────────

    def explore_schema(self) -> str:
        """Return a full schema summary: all tables, columns, types, and samples."""
        return _explore.explore_schema(self._session)

    def describe_table(self, table_name: str) -> str:
        """Return detailed column info, row count, and sample values for one table."""
        return _explore.describe_table(table_name, self._session)

    def list_tables(self) -> list[str]:
        """Return the names of all tables in the connected database."""
        return _query.list_tables(self._session)

    # ── Query ───────────────────────────────────────────────────────────────

    def run_query(self, sql: str) -> dict:
        """Execute a SELECT query and return results as a dict.

        Returns:
            {"columns": [...], "rows": [...], "row_count": N}

        Only SELECT / WITH statements are permitted. Results capped at 500 rows.
        Raises ValueError for disallowed queries.
        """
        return _query.run_query(sql, self._session)

    # ── Context manager ──────────────────────────────────────────────────────

    def __enter__(self):
        return self

    def __exit__(self, *_):
        self.disconnect()


__all__ = ["PocketAnalyst"]
