import os
import codecs
import re
import hashlib
import asyncio
import inspect
from concurrent.futures import ThreadPoolExecutor
from dotenv import load_dotenv

import requests
import redis
from .utils import Filter, convert_custom_filter
import json
from enum import Enum


from quillsql.db.cached_connection import CachedConnection
from quillsql.db.db_helper import (
    get_db_credentials,
    get_schema_column_info_by_db,
    get_schema_tables_by_db,
)
from quillsql.utils.schema_conversion import convert_type_to_postgres
from quillsql.utils.run_query_processes import (
    array_to_map,
    remove_fields,
)
from quillsql.utils.tenants import extract_tenant_ids
from quillsql.utils.pivot_template import (
    parse_distinct_values,
    hydrate_pivot_template,
)

load_dotenv()

ENV = os.getenv("PYTHON_ENV")
DEV_HOST = "http://localhost:8080"
PROD_HOST = "https://api.quill.co"
HOST = DEV_HOST if ENV == "development" else PROD_HOST

SINGLE_TENANT = "QUILL_SINGLE_TENANT"
ALL_TENANTS = "QUILL_ALL_TENANTS"
FLAG_TASKS = {'dashboard', 'report', 'item', 'report-info', 'filter-options'}

# Quill - Fullstack API Platform for Dashboards and Reporting.
class Quill:
    def __init__(
        self,
        private_key,
        database_type,
        database_connection_string=None,
        database_config=None,
        metadataServerURL=None,
        cache=None,
    ):
        if private_key is None:
            raise ValueError("Private key is required")
        if database_type is None:
            raise ValueError("Database type is required")
        if database_connection_string is None and database_config is None:
            raise ValueError("You must provide either DatabaseConnectionString or DatabaseConfig")

        # Handles both dsn-style connection strings (eg. "dbname=test password=secret" )
        # as well as url-style connection strings (eg. "postgres://foo@db.com")
        self.baseUrl = metadataServerURL if metadataServerURL is not None else HOST
        if database_connection_string is not None:
            self.target_connection = CachedConnection(
                database_type,
                get_db_credentials(database_type, database_connection_string),
                cache,
                True,
            )
        else:
            self.target_connection = CachedConnection(
                database_type, database_config, cache, False
            )
        self.private_key = private_key

    def close(self):
        connection = getattr(self, "target_connection", None)
        if connection and hasattr(connection, "close"):
            connection.close()

    async def aclose(self):
        connection = getattr(self, "target_connection", None)
        if connection and hasattr(connection, "close_async"):
            await connection.close_async()
            return
        self.close()

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()
        return False

    async def __aenter__(self):
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        await self.aclose()
        return False

    def __del__(self):
        try:
            self.close()
        except Exception:
            pass

    def get_cache(self, cache_config):
        cache_type = cache_config and cache_config.get("cache_type")
        if cache_type and cache_type == "redis" or cache_type == "rediss":
            return redis.Redis(
                host=cache_config.get("host", "localhost"),
                port=cache_config.get("port", 6379),
                username=cache_config.get("username", "default"),
                password=cache_config.get("password"),
            )
        return None

    def _get_tenant_mapped_flags_cache_key(self, client_id, tenants, report_id):
        if not client_id or not isinstance(tenants, list) or not tenants:
            return None
        if not isinstance(tenants[0], dict):
            return None

        normalized_tenants = sorted(
            [
                {
                    "tenantField": tenant.get("tenantField"),
                    "tenantIds": sorted(
                        list(tenant.get("tenantIds") or []), key=lambda value: str(value)
                    ),
                }
                for tenant in tenants
                if isinstance(tenant, dict)
            ],
            key=lambda tenant: tenant.get("tenantField") or "",
        )
        return (
            f"tenant-mapped-flags:{client_id}:{report_id or ''}:"
            f"{json.dumps(normalized_tenants, separators=(',', ':'), ensure_ascii=False)}"
        )

    def _get_tenant_mapped_flags_from_cache(self, cache_key, overwrite_cache):
        cache = getattr(self.target_connection, "cache", None)
        if overwrite_cache or not cache_key or not cache:
            return None
        try:
            cached_response = cache.get(cache_key)
            if not cached_response:
                return None
            parsed_response = json.loads(cached_response)
            is_valid_cache = (
                isinstance(parsed_response, list)
                and all(
                    isinstance(item, dict)
                    and isinstance(item.get("tenantField"), str)
                    and isinstance(item.get("flags"), list)
                    for item in parsed_response
                )
            )
            if not is_valid_cache:
                try:
                    cache.delete(cache_key)
                except Exception:
                    pass
                return None
            return parsed_response
        except Exception:
            try:
                if cache and cache_key:
                    cache.delete(cache_key)
            except Exception:
                pass
            return None

    def _set_tenant_mapped_flags_cache(self, cache_key, tenant_flags):
        cache = getattr(self.target_connection, "cache", None)
        if not cache_key or not cache:
            return
        ttl = getattr(self.target_connection, "ttl", None)
        try:
            if ttl:
                cache.set(cache_key, json.dumps(tenant_flags), "EX", ttl)
            else:
                cache.set(cache_key, json.dumps(tenant_flags))
        except Exception:
            pass

    def _post_quill_sync(self, path, payload, use_cache=True):
        if use_cache:
            result = self.post_quill(path, payload)
        else:
            result = self.post_quill(path, payload, use_cache)
        if inspect.isawaitable(result):
            try:
                asyncio.get_running_loop()
            except RuntimeError:
                return asyncio.run(result)
            with ThreadPoolExecutor(max_workers=1) as executor:
                return executor.submit(lambda: asyncio.run(result)).result()
        return result

    async def query(
        self,
        tenants,
        metadata,
        flags=None,
        filters: list[Filter] = None,
        admin_enabled: bool = None,
    ):
        if getattr(self.target_connection, "database_type", "").lower() in (
            "postgres",
            "postgresql",
        ):
            return await self._query_async(
                tenants, metadata, flags, filters, admin_enabled
            )
        return await asyncio.to_thread(
            self._query_sync,
            tenants,
            metadata,
            flags,
            filters,
            admin_enabled,
        )

    def query_sync(
        self,
        tenants,
        metadata,
        flags=None,
        filters: list[Filter] = None,
        admin_enabled: bool = None,
    ):
        return self._query_sync(tenants, metadata, flags, filters, admin_enabled)

    async def _query_async(
        self,
        tenants,
        metadata,
        flags=None,
        filters: list[Filter] = None,
        admin_enabled: bool = None,
    ):
        if not tenants:
            raise ValueError("You may not pass an empty tenants array.")

        responseMetadata = {}
        if not metadata:
            return {"error": "Missing metadata.", "status": "error", "data": {}}

        task = metadata.get("task")
        if not task:
            return {"error": "Missing task.", "status": "error", "data": {}}
        overwrite_cache = bool(metadata.get("overwriteCache", False))

        try:
            self.target_connection.tenant_ids = extract_tenant_ids(tenants)

            if task == "pivot-template":
                pivot_payload = {
                    **metadata,
                    "tenants": tenants,
                    "flags": flags,
                }
                if admin_enabled is not None:
                    pivot_payload["adminEnabled"] = admin_enabled
                if filters is not None:
                    pivot_payload["sdkFilters"] = [
                        convert_custom_filter(f) for f in filters
                    ]
                pivot_template_response = await self.post_quill(
                    "pivot-template",
                    pivot_payload,
                )

                if pivot_template_response.get("error"):
                    return {
                        "status": "error",
                        "error": pivot_template_response.get("error"),
                        "data": pivot_template_response.get("metadata") or {},
                    }

                template = pivot_template_response.get("metadata", {}).get("template")
                config = pivot_template_response.get("metadata", {}).get("config")
                distinct_values_query = pivot_template_response.get(
                    "metadata", {}
                ).get("distinctValuesQuery")
                row_count_query = pivot_template_response.get("metadata", {}).get(
                    "rowCountQuery"
                )

                distinct_values = []
                if distinct_values_query:
                    distinct_value_results = await self.run_queries_async(
                        [distinct_values_query],
                        self.target_connection.database_type,
                        metadata.get("databaseType"),
                        metadata,
                        None,
                        overwrite_cache,
                    )

                    distinct_value_query_results = (
                        distinct_value_results or {}
                    ).get("queryResults") or []
                    if distinct_value_query_results:
                        distinct_values = parse_distinct_values(
                            distinct_value_query_results[0], config.get("databaseType")
                        )

                try:
                    final_query = hydrate_pivot_template(template, distinct_values, config)
                except Exception as err:
                    return {
                        "status": "error",
                        "error": f"Failed to hydrate pivot template: {str(err)}",
                        "data": {},
                    }

                queries_to_run = [final_query]
                if row_count_query:
                    hydrated_row_count_query = hydrate_pivot_template(
                        row_count_query,
                        distinct_values,
                        config,
                    )
                    queries_to_run.append(hydrated_row_count_query)

                final_results = await self.run_queries_async(
                    queries_to_run,
                    self.target_connection.database_type,
                    metadata.get("databaseType"),
                    metadata,
                    pivot_template_response.get("metadata", {}).get("runQueryConfig"),
                    overwrite_cache,
                )

                responseMetadata = pivot_template_response.get("metadata") or {}
                if (
                    final_results.get("queryResults")
                    and len(final_results["queryResults"]) >= 1
                ):
                    query_results = final_results["queryResults"][0]
                    if query_results.get("rows"):
                        responseMetadata["rows"] = query_results["rows"]
                    if query_results.get("fields"):
                        responseMetadata["fields"] = query_results["fields"]

                if "template" in responseMetadata:
                    del responseMetadata["template"]
                if "distinctValuesQuery" in responseMetadata:
                    del responseMetadata["distinctValuesQuery"]
                if "rowCountQuery" in responseMetadata:
                    del responseMetadata["rowCountQuery"]

                return {
                    "data": responseMetadata,
                    "queries": final_results,
                    "status": "success",
                }

            tenant_flags = None
            if (
                task in FLAG_TASKS
                and tenants[0] != ALL_TENANTS
                and tenants[0] != SINGLE_TENANT
            ):
                tenant_mapped_flags_cache_key = self._get_tenant_mapped_flags_cache_key(
                    metadata.get("clientId"),
                    tenants,
                    metadata.get("reportId") or "",
                )
                cached_tenant_flags = self._get_tenant_mapped_flags_from_cache(
                    tenant_mapped_flags_cache_key,
                    overwrite_cache,
                )
                if cached_tenant_flags:
                    tenant_flags = cached_tenant_flags
                else:
                    tenant_flags_payload = {
                        "reportId": metadata.get("reportId")
                        or metadata.get("dashboardItemId"),
                        "clientId": metadata.get("clientId"),
                        "dashboardName": metadata.get("name"),
                        "tenants": tenants,
                        "flags": flags,
                    }
                    if admin_enabled is not None:
                        tenant_flags_payload["adminEnabled"] = admin_enabled
                    response = await self.post_quill(
                        "tenant-mapped-flags", tenant_flags_payload
                    )

                    if response.get("error"):
                        return {
                            "status": "error",
                            "error": response.get("error"),
                            "data": response.get("metadata") or {},
                        }

                    flag_query_results = await self.run_queries_async(
                        response.get("queries"),
                        self.target_connection.database_type,
                        overwrite_cache=overwrite_cache,
                    )

                    tenant_flags = []
                    query_order = (response.get("metadata") or {}).get("queryOrder") or []
                    query_results = (flag_query_results or {}).get("queryResults") or []
                    for index, tenant_field in enumerate(query_order):
                        query_result = (
                            query_results[index] if index < len(query_results) else {}
                        )
                        rows = (
                            query_result.get("rows")
                            if isinstance(query_result, dict)
                            else []
                        )
                        ordered_flags = []
                        seen = set()
                        for row in rows or []:
                            if not isinstance(row, dict):
                                continue
                            flag = row.get("quill_flag")
                            if flag is None or flag in seen:
                                continue
                            seen.add(flag)
                            ordered_flags.append(flag)
                        tenant_flags.append(
                            {
                                "tenantField": tenant_field,
                                "flags": ordered_flags,
                            }
                        )
                    self._set_tenant_mapped_flags_cache(
                        tenant_mapped_flags_cache_key,
                        tenant_flags,
                    )

            elif tenants[0] == SINGLE_TENANT and flags:
                if flags and isinstance(flags[0], dict):
                    tenant_flags = [{"tenantField": SINGLE_TENANT, "flags": flags}]
                else:
                    tenant_flags = flags

            if metadata.get("preQueries"):
                pre_query_results = await self.run_queries_async(
                    metadata.get("preQueries"),
                    self.target_connection.database_type,
                    metadata.get("databaseType"),
                    metadata,
                    metadata.get("runQueryConfig"),
                    overwrite_cache,
                )
            else:
                pre_query_results = {}

            if metadata.get("runQueryConfig") and metadata.get("runQueryConfig").get(
                "overridePost"
            ):
                return {"data": pre_query_results, "status": "success"}
            view_query = None
            if metadata.get("preQueries"):
                view_query = metadata.get("preQueries")[0]
            pre_query_columns = (
                pre_query_results.get("columns")
                if metadata.get("runQueryConfig")
                and metadata.get("runQueryConfig").get("getColumns")
                else None
            )
            payload = {
                **metadata,
                "tenants": tenants,
                "flags": tenant_flags,
                "viewQuery": view_query,
            }
            if pre_query_columns is not None:
                payload["preQueryResultsColumns"] = pre_query_columns
            if admin_enabled is not None:
                payload["adminEnabled"] = admin_enabled
            if filters is not None:
                payload["sdkFilters"] = [convert_custom_filter(f) for f in filters]
            quill_results = await self.post_quill(metadata.get("task"), payload)
            if quill_results.get("error"):
                responseMetadata = quill_results.get("metadata")
                response = {
                    "error": quill_results.get("error"),
                    "status": "error",
                    "data": {},
                }
                if responseMetadata:
                    response["data"] = responseMetadata
                return response

            if not quill_results.get("metadata"):
                quill_results["metadata"] = {}
            metadata = quill_results.get("metadata")
            responseMetadata = metadata
            results = await self.run_queries_async(
                quill_results.get("queries"),
                self.target_connection.database_type,
                metadata.get("databaseType"),
                metadata,
                metadata.get("runQueryConfig"),
                overwrite_cache,
            )

            should_wrap_results = isinstance(results, list) or not results
            if should_wrap_results:
                normalized_results = {
                    "queryResults": results if isinstance(results, list) else []
                }
            else:
                normalized_results = results

            if (
                should_wrap_results
                and not normalized_results.get("queryResults")
                and quill_results.get("queries")
            ):
                normalized_results["queryResults"] = (
                    normalized_results.get("queryResults") or []
                )

            if (
                normalized_results.get("mapped_array")
                and metadata.get("runQueryConfig", {}).get("arrayToMap")
            ):
                array_to_map = metadata["runQueryConfig"]["arrayToMap"]
                target_collection = responseMetadata.get(array_to_map["arrayName"])
                if isinstance(target_collection, list):
                    for index, array in enumerate(normalized_results["mapped_array"]):
                        if index >= len(target_collection):
                            continue
                        target_entry = target_collection[index]
                        if not isinstance(target_entry, dict):
                            target_entry = {}
                            target_collection[index] = target_entry
                        target_entry[array_to_map["field"]] = (
                            array if array is not None else []
                        )
                del normalized_results["mapped_array"]

            query_results_list = normalized_results.get("queryResults") or []
            if len(query_results_list) == 1 and isinstance(query_results_list[0], dict):
                query_result = query_results_list[0]
                quill_results["metadata"]["rows"] = query_result.get("rows")
                quill_results["metadata"]["fields"] = query_result.get("fields")
            return {
                "data": quill_results.get("metadata"),
                "queries": normalized_results,
                "status": "success",
            }

        except Exception as err:
            if task == "update-view":
                broken_view_payload = {
                    "table": metadata.get("name"),
                    "clientId": metadata.get("clientId"),
                    "error": str(err),
                }
                if admin_enabled is not None:
                    broken_view_payload["adminEnabled"] = admin_enabled
                await self.post_quill("set-broken-view", broken_view_payload)
            return {
                "error": str(err).splitlines()[0],
                "status": "error",
                "data": responseMetadata,
            }

    def _query_sync(
        self,
        tenants,
        metadata,
        flags=None,
        filters: list[Filter] = None,
        admin_enabled: bool = None,
    ):
        if not tenants:
            raise ValueError("You may not pass an empty tenants array.")

        responseMetadata = {}
        if not metadata:
            return {"error": "Missing metadata.", "status": "error", "data": {}}

        task = metadata.get("task")
        if not task:
            return {"error": "Missing task.", "status": "error", "data": {}}
        overwrite_cache = bool(metadata.get("overwriteCache", False))

        try:
            # Set tenant IDs in the connection
            self.target_connection.tenant_ids = extract_tenant_ids(tenants)

            # Handle pivot-template task
            if task == "pivot-template":
                pivot_payload = {
                    **metadata,
                    "tenants": tenants,
                    "flags": flags,
                }
                if admin_enabled is not None:
                    pivot_payload["adminEnabled"] = admin_enabled
                if filters is not None:
                    pivot_payload["sdkFilters"] = [
                        convert_custom_filter(f) for f in filters
                    ]
                pivot_template_response = self._post_quill_sync(
                    "pivot-template",
                    pivot_payload,
                )

                if pivot_template_response.get("error"):
                    return {
                        "status": "error",
                        "error": pivot_template_response.get("error"),
                        "data": pivot_template_response.get("metadata") or {},
                    }

                template = pivot_template_response.get("metadata", {}).get("template")
                config = pivot_template_response.get("metadata", {}).get("config")
                distinct_values_query = pivot_template_response.get("metadata", {}).get("distinctValuesQuery")
                row_count_query = pivot_template_response.get("metadata", {}).get("rowCountQuery")

                distinct_values = []
                if distinct_values_query:
                    distinct_value_results = self.run_queries(
                        [distinct_values_query],
                        self.target_connection.database_type,
                        metadata.get("databaseType"),
                        metadata,
                        None,
                        overwrite_cache,
                    )

                    distinct_value_query_results = (
                        distinct_value_results or {}
                    ).get("queryResults") or []
                    if distinct_value_query_results:
                        distinct_values = parse_distinct_values(
                            distinct_value_query_results[0],
                            config.get("databaseType")
                        )

                try:
                    final_query = hydrate_pivot_template(template, distinct_values, config)
                except Exception as err:
                    return {
                        "status": "error",
                        "error": f"Failed to hydrate pivot template: {str(err)}",
                        "data": {},
                    }

                queries_to_run = [final_query]
                if row_count_query:
                    hydrated_row_count_query = hydrate_pivot_template(
                        row_count_query,
                        distinct_values,
                        config
                    )
                    queries_to_run.append(hydrated_row_count_query)

                final_results = self.run_queries(
                    queries_to_run,
                    self.target_connection.database_type,
                    metadata.get("databaseType"),
                    metadata,
                    pivot_template_response.get("metadata", {}).get("runQueryConfig"),
                    overwrite_cache,
                )

                responseMetadata = pivot_template_response.get("metadata") or {}
                if final_results.get("queryResults") and len(final_results["queryResults"]) >= 1:
                    query_results = final_results["queryResults"][0]
                    if query_results.get("rows"):
                        responseMetadata["rows"] = query_results["rows"]
                    if query_results.get("fields"):
                        responseMetadata["fields"] = query_results["fields"]

                if "template" in responseMetadata:
                    del responseMetadata["template"]
                if "distinctValuesQuery" in responseMetadata:
                    del responseMetadata["distinctValuesQuery"]
                if "rowCountQuery" in responseMetadata:
                    del responseMetadata["rowCountQuery"]

                return {
                    "data": responseMetadata,
                    "queries": final_results,
                    "status": "success",
                }

            # Handle tenant flags synthesis
            tenant_flags = None
            if (task in FLAG_TASKS and 
                tenants[0] != ALL_TENANTS and 
                tenants[0] != SINGLE_TENANT
            ):
                tenant_mapped_flags_cache_key = self._get_tenant_mapped_flags_cache_key(
                    metadata.get("clientId"),
                    tenants,
                    metadata.get("reportId") or "",
                )
                cached_tenant_flags = self._get_tenant_mapped_flags_from_cache(
                    tenant_mapped_flags_cache_key,
                    overwrite_cache,
                )
                if cached_tenant_flags:
                    tenant_flags = cached_tenant_flags
                else:
                    tenant_flags_payload = {
                        'reportId': metadata.get('reportId') or metadata.get('dashboardItemId'),
                        'clientId': metadata.get('clientId'),
                        'dashboardName': metadata.get('name'),
                        'tenants': tenants,
                        'flags': flags,
                    }
                    if admin_enabled is not None:
                        tenant_flags_payload['adminEnabled'] = admin_enabled
                    response = self._post_quill_sync('tenant-mapped-flags', tenant_flags_payload)
                    
                    if response.get('error'):
                        return {
                            'status': 'error',
                            'error': response.get('error'),
                            'data': response.get('metadata') or {},
                        }
                    
                    flag_query_results = self.run_queries(
                        response.get('queries'),
                        self.target_connection.database_type,
                        overwrite_cache=overwrite_cache,
                    )
                    
                    tenant_flags = []
                    query_order = (response.get('metadata') or {}).get('queryOrder') or []
                    query_results = (flag_query_results or {}).get('queryResults') or []
                    for index, tenant_field in enumerate(query_order):
                        query_result = (
                            query_results[index]
                            if index < len(query_results)
                            else {}
                        )
                        rows = (
                            query_result.get('rows')
                            if isinstance(query_result, dict)
                            else []
                        )
                        ordered_flags = []
                        seen = set()
                        for row in rows or []:
                            if not isinstance(row, dict):
                                continue
                            flag = row.get('quill_flag')
                            if flag is None or flag in seen:
                                continue
                            seen.add(flag)
                            ordered_flags.append(flag)
                        tenant_flags.append({
                            'tenantField': tenant_field,
                            'flags': ordered_flags,
                        })
                    self._set_tenant_mapped_flags_cache(
                        tenant_mapped_flags_cache_key,
                        tenant_flags,
                    )

            elif tenants[0] == SINGLE_TENANT and flags:
                if flags and isinstance(flags[0], dict):
                    tenant_flags = [{'tenantField': SINGLE_TENANT, 'flags': flags}]
                else:
                    tenant_flags = flags

            if metadata.get("preQueries"):
                pre_query_results = self.run_queries(
                    metadata.get("preQueries"),
                    self.target_connection.database_type,
                    metadata.get("databaseType"),
                    metadata,
                    metadata.get("runQueryConfig"),
                    overwrite_cache,
                )
            else:
                pre_query_results = {}

            if metadata.get("runQueryConfig") and metadata.get("runQueryConfig").get(
                "overridePost"
            ):
                return {"data": pre_query_results, "status": "success"}
            view_query = None
            if metadata.get("preQueries"):
                view_query = metadata.get("preQueries")[0]
            pre_query_columns = (
                pre_query_results.get("columns")
                if metadata.get("runQueryConfig")
                and metadata.get("runQueryConfig").get("getColumns")
                else None
            )
            payload = {
                **metadata,
                "tenants": tenants,
                "flags": tenant_flags,
                "viewQuery": view_query,
            }
            if pre_query_columns is not None:
                payload["preQueryResultsColumns"] = pre_query_columns
            if admin_enabled is not None:
                payload["adminEnabled"] = admin_enabled
            if filters is not None:
                payload["sdkFilters"] = [convert_custom_filter(f) for f in filters]
            quill_results = self._post_quill_sync(metadata.get("task"), payload)
            if quill_results.get("error"):
                responseMetadata = quill_results.get("metadata")
                response = {
                    "error": quill_results.get("error"),
                    "status": "error",
                    "data": {},
                }
                if responseMetadata:
                    response["data"] = responseMetadata
                return response

            # If there is no metadata in the quill results, create one
            if not quill_results.get("metadata"):
                quill_results["metadata"] = {}
            metadata = quill_results.get("metadata")
            responseMetadata = metadata
            results = self.run_queries(
                quill_results.get("queries"),
                self.target_connection.database_type,
                metadata.get("databaseType"),
                metadata,
                metadata.get("runQueryConfig"),
                overwrite_cache,
            )

            should_wrap_results = isinstance(results, list) or not results
            if should_wrap_results:
                normalized_results = {
                    "queryResults": results if isinstance(results, list) else []
                }
            else:
                normalized_results = results

            if (
                should_wrap_results
                and not normalized_results.get("queryResults")
                and quill_results.get("queries")
            ):
                normalized_results["queryResults"] = (
                    normalized_results.get("queryResults") or []
                )

            if (
                normalized_results.get("mapped_array")
                and metadata.get("runQueryConfig", {}).get("arrayToMap")
            ):
                array_to_map = metadata["runQueryConfig"]["arrayToMap"]
                target_collection = responseMetadata.get(array_to_map["arrayName"])
                if isinstance(target_collection, list):
                    for index, array in enumerate(normalized_results["mapped_array"]):
                        if index >= len(target_collection):
                            continue
                        target_entry = target_collection[index]
                        if not isinstance(target_entry, dict):
                            target_entry = {}
                            target_collection[index] = target_entry
                        target_entry[array_to_map["field"]] = array if array is not None else []
                del normalized_results["mapped_array"]

            query_results_list = normalized_results.get("queryResults") or []
            if len(query_results_list) == 1 and isinstance(query_results_list[0], dict):
                query_result = query_results_list[0]
                quill_results["metadata"]["rows"] = query_result.get("rows")
                quill_results["metadata"]["fields"] = query_result.get("fields")
            return {
                "data": quill_results.get("metadata"),
                "queries": normalized_results,
                "status": "success",
            }

        except Exception as err:
            if task == "update-view":
                broken_view_payload = {
                    "table": metadata.get("name"),
                    "clientId": metadata.get("clientId"),
                    "error": str(err),
                }
                if admin_enabled is not None:
                    broken_view_payload["adminEnabled"] = admin_enabled
                self._post_quill_sync("set-broken-view", broken_view_payload)
            return {
                "error": str(err).splitlines()[0],
                "status": "error",
                "data": responseMetadata,
            }
        
    async def stream(
        self,
        tenants,
        metadata,
        flags=None,
        filters=None,
        admin_enabled=None,
    ):
        if not tenants:
            raise ValueError("You may not pass an empty tenants array.")

        if not metadata:
            yield {"type": "error", "errorText": "Missing metadata."}
            return

        task = metadata.get("task")
        if not task:
            yield {"type": "error", "errorText": "Missing task."}
            return

        try:
            # Set tenant IDs in the connection
            self.target_connection.tenant_ids = extract_tenant_ids(tenants)
            if task in ("chat", "agent"):
                for event in self._agentic_chat_loop(
                    tenants,
                    metadata,
                    flags,
                    filters,
                    admin_enabled,
                ):
                    yield event
                return

            for event in self._stream_sse(
                task,
                tenants,
                metadata,
                flags,
                filters,
                admin_enabled,
            ):
                yield event
            return
        except Exception as err:
            yield {
                "type": "error",
                "errorText": str(err).splitlines()[0],
            }
            return

    def _normalize_tenant_flags(self, tenants, flags):
        tenant_flags = None
        if tenants and tenants[0] == SINGLE_TENANT and flags:
            if flags and isinstance(flags[0], dict):
                tenant_flags = [{"tenantField": SINGLE_TENANT, "flags": flags}]
            else:
                tenant_flags = flags
        return tenant_flags

    def _agentic_chat_loop(self, tenants, metadata, flags, filters, admin_enabled):
        messages = list(metadata.get("messages") or [])
        max_iterations = 10

        for _ in range(max_iterations):
            payload = {
                **metadata,
                "messages": messages,
            }

            has_tool_calls = False
            assistant_text = ""
            tool_results = []
            tool_calls = []

            for event in self._stream_sse(
                "agent",
                tenants,
                payload,
                flags,
                filters,
                admin_enabled,
            ):
                yield event

                if event.get("type") == "text-delta":
                    assistant_text += event.get("delta", "")

                if event.get("type") == "tool-input-available":
                    tool_name = event.get("toolName")
                    tool_call_id = event.get("toolCallId")
                    tool_input = event.get("input") or {}

                    if tool_call_id is None:
                        yield {
                            "type": "error",
                            "errorText": "Missing toolCallId for tool-input-available event.",
                        }
                        continue

                    has_tool_calls = True
                    yield {
                        "type": "tool-executing",
                        "toolCallId": tool_call_id,
                        "toolName": tool_name,
                    }

                    result = self._execute_tool_locally(tool_name, tool_input)
                    tool_results.append(
                        {
                            "toolCallId": tool_call_id,
                            "toolName": tool_name,
                            "input": tool_input,
                            "result": result,
                        }
                    )
                    tool_calls.append(
                        {
                            "id": tool_call_id,
                            "type": "function",
                            "function": {
                                "name": tool_name,
                                "arguments": json.dumps(tool_input or {}),
                            },
                        }
                    )

                    yield {
                        "type": "tool-result",
                        "toolCallId": tool_call_id,
                        "result": result,
                    }

                if event.get("type") in ("finish", "error"):
                    break

            if not has_tool_calls:
                break

            def _build_tool_result_for_history(tool_result):
                result = tool_result.get("result") or {}
                has_rows = isinstance(result.get("rows"), list)
                has_fields = isinstance(result.get("fields"), list)
                is_query_result = has_rows or has_fields or result.get("dbMismatched")
                if not is_query_result:
                    return result
                tool_input = tool_result.get("input") or {}
                error = result.get("error") or tool_input.get("error")
                status = "error" if error or result.get("dbMismatched") else "success"
                payload = {"status": status}
                if tool_input.get("sql"):
                    payload["sql"] = tool_input.get("sql")
                if error:
                    payload["error"] = error
                if result.get("dbMismatched"):
                    payload["meta"] = {
                        "dbMismatched": True,
                        "backendDatabaseType": result.get("backendDatabaseType"),
                    }
                elif not error:
                    payload["meta"] = {
                        "rowsFetchedSuccessfully": True,
                        "rowCount": len(result.get("rows") or []),
                    }
                return payload

            messages.append(
                {
                    "role": "assistant",
                    "content": assistant_text or None,
                    "tool_calls": tool_calls,
                }
            )
            for tool_result in tool_results:
                messages.append(
                    {
                        "role": "tool",
                        "tool_call_id": tool_result["toolCallId"],
                        "content": json.dumps(_build_tool_result_for_history(tool_result)),
                    }
                )

        yield {"type": "done"}

    def _execute_tool_locally(self, tool_name, tool_input):
        if tool_name == "generateReport":
            if tool_input.get("error"):
                return {"error": tool_input.get("error")}
            sql = tool_input.get("sql")
            if not sql:
                return {"error": "No SQL provided"}
            results = self.run_queries(
                [sql],
                self.target_connection.database_type,
            )
            if results.get("dbMismatched"):
                return results
            query_results = results.get("queryResults") or []
            if query_results and isinstance(query_results[0], dict):
                if query_results[0].get("error"):
                    return query_results[0]
                return {
                    "rows": query_results[0].get("rows", []),
                    "fields": query_results[0].get("fields", []),
                }
            return {"rows": [], "fields": []}

        if tool_name == "createChart":
            return {"chartConfig": tool_input}

        return {"error": f"Unknown tool: {tool_name}"}

    def _stream_sse(self, endpoint, tenants, payload, flags, filters, admin_enabled):
        tenant_flags = self._normalize_tenant_flags(tenants, flags)
        request_payload = {
            **payload,
            "tenants": tenants,
            "flags": tenant_flags,
        }
        if filters:
            request_payload["sdkFilters"] = [convert_custom_filter(f) for f in filters]
        if admin_enabled is not None:
            request_payload["adminEnabled"] = admin_enabled

        # Custom JSON Encoder to handle Enums
        class EnumEncoder(json.JSONEncoder):
            def default(self, obj):
                if isinstance(obj, Enum):
                    return obj.value  # Convert enum to its value (string in this case)
                return super().default(obj)

        url = f"{self.baseUrl}/sdk/{endpoint}"
        headers = {
            "Authorization": f"Bearer {self.private_key}",
            "Content-Type": "application/json",
            "Accept": "text/event-stream",
        }
        encoded = json.dumps(request_payload, cls=EnumEncoder)

        resp = requests.post(url, data=encoded, headers=headers, stream=True)
        decoder = codecs.getincrementaldecoder("utf-8")()
        buf = ""
        for chunk in resp.iter_content(chunk_size=4096):
            buf += decoder.decode(chunk)
            while "\n\n" in buf:
                raw_event, buf = buf.split("\n\n", 1)
                data_lines = []
                for line in raw_event.splitlines():
                    if line.startswith("data:"):
                        data_lines.append(line[len("data:"):].strip())
                if not data_lines:
                    continue
                payload = "\n".join(data_lines)
                if payload == "[DONE]":
                    return
                try:
                    parsed = json.loads(payload)
                    if isinstance(parsed, str):
                        yield {"type": "text-delta", "id": "0", "delta": parsed}
                    else:
                        yield parsed
                except json.JSONDecodeError:
                    continue

        # flush any partial code points at the end
        buf += decoder.decode(b"", final=True)
        yield buf
        return

    def apply_limit(self, query, limit):
        # Simple logic: if query already has a limit, don't add another
        if getattr(self.target_connection, 'database_type', '').lower() == 'mssql':
            import re
            if re.search(r'SELECT TOP \\d+', query, re.IGNORECASE):
                return query
            return re.sub(r'select', f'SELECT TOP {limit}', query, flags=re.IGNORECASE)
        else:
            if 'limit ' in query.lower():
                return query
            return f"{query.rstrip(';')} limit {limit}"

    async def run_queries_async(
        self,
        queries,
        pkDatabaseType,
        databaseType=None,
        metadata=None,
        runQueryConfig=None,
        overwrite_cache=False,
    ):
        if getattr(self.target_connection, "database_type", "").lower() not in (
            "postgres",
            "postgresql",
        ):
            return self.run_queries(
                queries,
                pkDatabaseType,
                databaseType,
                metadata,
                runQueryConfig,
                overwrite_cache,
            )

        results = {}
        if not queries:
            return {"queryResults": []}

        def _normalize_db_type(db_type):
            if not db_type:
                return None
            lowered = db_type.lower()
            if lowered in ("postgresql", "postgres"):
                return "postgres"
            return lowered

        normalized_metadata_db = _normalize_db_type(databaseType)
        normalized_backend_db = _normalize_db_type(pkDatabaseType)

        def _is_read_only_query(query):
            if not isinstance(query, str):
                return False
            normalized = query.strip().lower()
            if not normalized:
                return False
            if not normalized.startswith(("select", "with", "show", "explain")):
                return False
            if "for update" in normalized or "for share" in normalized:
                return False
            if normalized.startswith("with") and re.search(
                r"\b(insert|update|delete|merge)\b", normalized
            ):
                return False
            return not re.search(
                r"\b(insert|update|delete|merge|alter|drop|create|truncate|grant|revoke|lock|vacuum|analyze|refresh)\b",
                normalized,
            )

        def _can_parallelize_query_batch(query_batch):
            return (
                normalized_backend_db == "postgres"
                and len(query_batch) > 1
                and all(_is_read_only_query(query) for query in query_batch)
            )

        def _get_max_concurrency():
            configured_max = (runQueryConfig or {}).get("maxConcurrency", 8)
            if not isinstance(configured_max, int) or configured_max < 1:
                configured_max = 8
            return min(len(queries), configured_max)

        async def _query_many(query_batch):
            semaphore = asyncio.Semaphore(_get_max_concurrency())

            async def _run(query):
                async with semaphore:
                    return await self.target_connection.query_async(
                        query, overwrite_cache
                    )

            return await asyncio.gather(*(_run(query) for query in query_batch))

        if (
            normalized_metadata_db
            and normalized_backend_db
            and normalized_metadata_db != normalized_backend_db
        ):
            return {"dbMismatched": True, "backendDatabaseType": pkDatabaseType}
        if runQueryConfig and runQueryConfig.get("arrayToMap"):
            mapped_array = array_to_map(
                queries,
                runQueryConfig.get("arrayToMap"),
                metadata,
                self.target_connection,
            )

            return {"queryResults": [], "mapped_array": mapped_array}
        elif runQueryConfig and runQueryConfig.get("getColumns"):
            query_results = await self.target_connection.query_async(
                queries[0].strip().rstrip(";") + " limit 1000",
                overwrite_cache,
            )
            results["columns"] = [
                {
                    "fieldType": convert_type_to_postgres(result["dataTypeID"]),
                    "name": result["name"],
                    "displayName": result["name"],
                    "isVisible": True,
                    "field": result["name"],
                }
                for result in query_results["fields"]
            ]
        elif runQueryConfig and runQueryConfig.get("getColumnsForSchema"):
            query_results = []
            for table in queries:
                if not table.get("viewQuery") or (
                    not table.get("isSelectStar") and not table.get("customFieldInfo")
                ):
                    query_results.append(table)
                    continue

                limit = ""
                if runQueryConfig.get("limitBy"):
                    limit = f" limit {runQueryConfig.get('limitBy')}"

                try:
                    query_result = await self.target_connection.query_async(
                        f"{table['viewQuery'].strip().rstrip(';')} {limit}",
                        overwrite_cache,
                    )
                    columns = [
                        {
                            "fieldType": convert_type_to_postgres(field["dataTypeID"]),
                            "name": field["name"],
                            "displayName": field["name"],
                            "isVisible": True,
                            "field": field["name"],
                        }
                        for field in query_result["fields"]
                    ]
                    query_results.append(
                        {**table, "columns": columns, "rows": query_result["rows"]}
                    )
                except Exception as e:
                    query_results.append(
                        {**table, "error": f"Error fetching columns {e}"}
                    )

            results["queryResults"] = query_results
            if runQueryConfig.get("fieldsToRemove"):
                results["queryResults"] = [
                    {
                        **table,
                        "columns": [
                            column
                            for column in table.get("columns", [])
                            if column["name"] not in runQueryConfig["fieldsToRemove"]
                        ],
                    }
                    for table in query_results
                ]
            return results
        elif runQueryConfig and runQueryConfig.get("getTables"):
            # Schema introspection path remains on the sync helper currently.
            tables = get_schema_tables_by_db(
                self.target_connection.database_type,
                self.target_connection.connection,
                runQueryConfig["schemaNames"],
            )
            schema = get_schema_column_info_by_db(
                self.target_connection.database_type,
                self.target_connection.connection,
                runQueryConfig["schemaNames"],
                tables,
            )
            results["queryResults"] = schema
        elif runQueryConfig and runQueryConfig.get("runIndividualQueries"):
            async def _run_individual_query(query):
                try:
                    run_query = query
                    if runQueryConfig.get("limitBy"):
                        run_query = self.apply_limit(query, runQueryConfig["limitBy"])
                    return await self.target_connection.query_async(
                        run_query, overwrite_cache
                    )
                except Exception as e:
                    return {
                        "query": query,
                        "error": str(e),
                    }

            if _can_parallelize_query_batch(queries):
                query_results = await asyncio.gather(
                    *(_run_individual_query(query) for query in queries)
                )
            else:
                query_results = []
                for query in queries:
                    query_results.append(await _run_individual_query(query))
            results["queryResults"] = query_results
        else:
            if runQueryConfig and runQueryConfig.get("limitThousand"):
                queries = [
                    query.strip().rstrip(";") + " limit 1000" for query in queries
                ]
            elif runQueryConfig and runQueryConfig.get("limitBy"):
                queries = [
                    query.strip().rstrip(";")
                    + f" limit {runQueryConfig.get('limitBy')}"
                    for query in queries
                ]
            if _can_parallelize_query_batch(queries):
                query_results = await _query_many(queries)
            else:
                query_results = []
                for query in queries:
                    query_results.append(
                        await self.target_connection.query_async(query, overwrite_cache)
                    )
            results["queryResults"] = query_results
            if runQueryConfig and runQueryConfig.get("fieldsToRemove"):
                results["queryResults"] = [
                    remove_fields(query_result, runQueryConfig.get("fieldsToRemove"))
                    for query_result in results["queryResults"]
                ]
            if runQueryConfig and runQueryConfig.get("convertDatatypes"):
                for query_result in results["queryResults"]:
                    query_result["fields"] = [
                        {
                            "name": field["name"],
                            "displayName": field["name"],
                            "field": field["name"],
                            "isVisible": True,
                            "dataTypeID": field["dataTypeID"],
                            "fieldType": convert_type_to_postgres(field["dataTypeID"]),
                        }
                        for field in query_result["fields"]
                    ]

        return results

    def run_queries(
        self,
        queries,
        pkDatabaseType,
        databaseType=None,
        metadata=None,
        runQueryConfig=None,
        overwrite_cache=False,
    ):
        results = {}
        if not queries:
            return {"queryResults": []}
        def _normalize_db_type(db_type):
            if not db_type:
                return None
            lowered = db_type.lower()
            if lowered in ("postgresql", "postgres"):
                return "postgres"
            return lowered

        normalized_metadata_db = _normalize_db_type(databaseType)
        normalized_backend_db = _normalize_db_type(pkDatabaseType)

        def _is_read_only_query(query):
            if not isinstance(query, str):
                return False
            normalized = query.strip().lower()
            if not normalized:
                return False
            if not normalized.startswith(("select", "with", "show", "explain")):
                return False
            if "for update" in normalized or "for share" in normalized:
                return False
            if normalized.startswith("with") and re.search(
                r"\b(insert|update|delete|merge)\b", normalized
            ):
                return False
            return not re.search(
                r"\b(insert|update|delete|merge|alter|drop|create|truncate|grant|revoke|lock|vacuum|analyze|refresh)\b",
                normalized,
            )

        def _can_parallelize_query_batch(query_batch):
            return (
                normalized_backend_db == "postgres"
                and len(query_batch) > 1
                and all(_is_read_only_query(query) for query in query_batch)
            )

        def _get_max_concurrency():
            configured_max = (runQueryConfig or {}).get("maxConcurrency", 8)
            if not isinstance(configured_max, int) or configured_max < 1:
                configured_max = 8
            return min(len(queries), configured_max)

        def _parallel_query_batch(query_batch):
            worker_count = _get_max_concurrency()
            with ThreadPoolExecutor(max_workers=worker_count) as executor:
                return list(
                    executor.map(
                        lambda query: self.target_connection.query(query, overwrite_cache),
                        query_batch,
                    )
                )

        if (
            normalized_metadata_db
            and normalized_backend_db
            and normalized_metadata_db != normalized_backend_db
        ):
            return {"dbMismatched": True, "backendDatabaseType": pkDatabaseType}
        if runQueryConfig and runQueryConfig.get("arrayToMap"):
            mapped_array = array_to_map(
                queries,
                runQueryConfig.get("arrayToMap"),
                metadata,
                self.target_connection,
            )

            return {"queryResults": [], "mapped_array": mapped_array}
        elif runQueryConfig and runQueryConfig.get("getColumns"):
            query_results = self.target_connection.query(
                queries[0].strip().rstrip(";") + " limit 1000",
                overwrite_cache,
            )
            results["columns"] = [
                {
                    "fieldType": convert_type_to_postgres(result["dataTypeID"]),
                    "name": result["name"],
                    "displayName": result["name"],
                    "isVisible": True,
                    "field": result["name"],
                }
                for result in query_results["fields"]
            ]
        elif runQueryConfig and runQueryConfig.get("getColumnsForSchema"):
            query_results = []
            for table in queries:
                if not table.get("viewQuery") or (
                    not table.get("isSelectStar") and not table.get("customFieldInfo")
                ):
                    query_results.append(table)
                    continue

                limit = ""
                if runQueryConfig.get("limitBy"):
                    limit = f" limit {runQueryConfig.get('limitBy')}"

                try:
                    query_result = self.target_connection.query(
                        f"{table['viewQuery'].strip().rstrip(';')} {limit}",
                        overwrite_cache,
                    )
                    columns = [
                        {
                            "fieldType": convert_type_to_postgres(field["dataTypeID"]),
                            "name": field["name"],
                            "displayName": field["name"],
                            "isVisible": True,
                            "field": field["name"],
                        }
                        for field in query_result["fields"]
                    ]
                    query_results.append(
                        {**table, "columns": columns, "rows": query_result["rows"]}
                    )
                except Exception as e:
                    query_results.append(
                        {**table, "error": f"Error fetching columns {e}"}
                    )

            results["queryResults"] = query_results
            if runQueryConfig.get("fieldsToRemove"):
                results["queryResults"] = [
                    {
                        **table,
                        "columns": [
                            column
                            for column in table.get("columns", [])
                            if column["name"] not in runQueryConfig["fieldsToRemove"]
                        ],
                    }
                    for table in query_results
                ]
            return results
        elif runQueryConfig and runQueryConfig.get("getTables"):
            tables = get_schema_tables_by_db(
                self.target_connection.database_type,
                self.target_connection.connection,
                runQueryConfig["schemaNames"],
            )
            schema = get_schema_column_info_by_db(
                self.target_connection.database_type,
                self.target_connection.connection,
                runQueryConfig["schemaNames"],
                tables,
            )
            results["queryResults"] = schema
        elif runQueryConfig and runQueryConfig.get("runIndividualQueries"):
            # so that one query doesn't fail the whole thing
            # the only reason this isn't the default behavior is for backwards compatibility
            def _run_individual_query(query):
                try:
                    run_query = query
                    if runQueryConfig.get("limitBy"):
                        run_query = self.apply_limit(query, runQueryConfig["limitBy"])
                    return self.target_connection.query(run_query, overwrite_cache)
                except Exception as e:
                    return {
                        "query": query,
                        "error": str(e),
                    }

            if _can_parallelize_query_batch(queries):
                worker_count = _get_max_concurrency()
                with ThreadPoolExecutor(max_workers=worker_count) as executor:
                    query_results = list(executor.map(_run_individual_query, queries))
            else:
                query_results = [_run_individual_query(query) for query in queries]
            results["queryResults"] = query_results
        else:
            if runQueryConfig and runQueryConfig.get("limitThousand"):
                queries = [
                    query.strip().rstrip(";") + " limit 1000" for query in queries
                ]
            elif runQueryConfig and runQueryConfig.get("limitBy"):
                queries = [
                    query.strip().rstrip(";")
                    + f" limit {runQueryConfig.get('limitBy')}"
                    for query in queries
                ]
            if _can_parallelize_query_batch(queries):
                query_results = _parallel_query_batch(queries)
            else:
                query_results = [
                    self.target_connection.query(query, overwrite_cache)
                    for query in queries
                ]
            results["queryResults"] = query_results
            if runQueryConfig and runQueryConfig.get("fieldsToRemove"):
                results["queryResults"] = [
                    remove_fields(query_result, runQueryConfig.get("fieldsToRemove"))
                    for query_result in results["queryResults"]
                ]
            if runQueryConfig and runQueryConfig.get("convertDatatypes"):
                for query_result in results["queryResults"]:
                    query_result["fields"] = [
                        {
                            "name": field["name"],
                            "displayName": field["name"],
                            "field": field["name"],
                            "isVisible": True,
                            "dataTypeID": field["dataTypeID"],
                            "fieldType": convert_type_to_postgres(field["dataTypeID"]),
                        }
                        for field in query_result["fields"]
                    ]

        return results
    
    async def post_quill(self, path, payload, use_cache=True):
        # Custom JSON Encoder to handle Enums
        class EnumEncoder(json.JSONEncoder):
            def default(self, obj):
                if isinstance(obj, Enum):
                    return obj.value  # Convert enum to its value (string in this case)
                return super().default(obj)

        cache = getattr(self.target_connection, "cache", None)
        # Match Node SDK behavior: skip post cache for tenant-mapped-flags and set-broken-view.
        can_use_cache = bool(
            use_cache and cache and path not in ("tenant-mapped-flags", "set-broken-view")
        )
        cache_key = None
        if can_use_cache:
            remove_cache_override = dict(payload or {})
            remove_cache_override.pop("overwriteCache", None)
            key_payload = json.dumps(
                remove_cache_override,
                cls=EnumEncoder,
                separators=(",", ":"),
                ensure_ascii=False,
            )
            hash_value = hashlib.sha256(key_payload.encode("utf-8")).hexdigest()
            cache_key = f"quill:post:{path}:{hash_value}"
            if not (payload or {}).get("overwriteCache"):
                try:
                    cached = cache.get(cache_key)
                    if cached:
                        try:
                            return json.loads(cached)
                        except Exception:
                            try:
                                cache.delete(cache_key)
                            except Exception:
                                pass
                except Exception:
                    pass

        url = f"{self.baseUrl}/sdk/{path}"
        # Set content type to application/json
        headers = {"Authorization": f"Bearer {self.private_key}", "Content-Type": "application/json"}
        encoded = json.dumps(payload, cls=EnumEncoder)
        response = await asyncio.to_thread(
            requests.post,
            url,
            data=encoded,
            headers=headers,
        )
        response_json = response.json()

        if can_use_cache and cache_key and not response_json.get("error"):
            ttl = getattr(self.target_connection, "ttl", None)
            try:
                if ttl:
                    cache.set(cache_key, json.dumps(response_json), "EX", ttl)
                else:
                    cache.set(cache_key, json.dumps(response_json))
            except Exception:
                pass

        return response_json
