import os

from backend.database.store import SQLiteStore
from dataiku.connector import Connector


class AgentHubTableConnector(Connector):
    def __init__(self, config, plugin_config):
        Connector.__init__(self, config, plugin_config)

        project_key = config.get("project_key")
        webapp_id = config.get("webapp_id")
        self.table_name = config.get("table_name")
        self.decompress = config.get("decompress_blobs", True)

        if not project_key:
            raise ValueError("Project key is required")
        if not webapp_id:
            raise ValueError("Webapp ID is required")
        if not self.table_name:
            raise ValueError("Table name is required")

        dip_home = os.environ.get("DIP_HOME")
        if not dip_home:
            raise ValueError("DIP_HOME environment variable not found")

        workload_folder = os.path.join(dip_home, "workload-folders", "webapps", project_key, webapp_id)

        if not os.path.exists(workload_folder):
            raise ValueError(f"Webapp workload folder not found: {workload_folder}")

        db_path = os.path.join(workload_folder, "data_store.db")
        if not os.path.exists(db_path):
            raise ValueError(f"Database file not found: {db_path}")

        self.store = SQLiteStore(db_path, max_artifacts_size_mb=None)

    def get_read_schema(self):
        schema_columns = self.store.get_table_schema(self.table_name)
        return {"columns": schema_columns}

    def generate_rows(self, dataset_schema=None, dataset_partitioning=None, partition_id=None, records_limit=-1):
        for row in self.store.export_table_rows(self.table_name, self.decompress):
            yield row
