# Copyright (c) 2025 Patrick Motsch
# All rights reserved.
"""Feature Data Sub-Agent tool (queryFeatureInstance)."""

import hashlib
import logging
import time
from typing import Any, Dict, List, Optional, Tuple

from modules.serviceCenter.services.serviceAgent.datamodelAgent import ToolResult
from modules.serviceCenter.services.serviceAgent.toolRegistry import ToolRegistry

from modules.serviceCenter.services.serviceAgent.coreTools._helpers import (
    _getOrCreateTempFolder,
    _looksLikeBinary,
    _resolveFileScope,
    _MAX_TOOL_RESULT_CHARS,
)

logger = logging.getLogger(__name__)

_featureDbConnPool: Dict[str, Any] = {}
_featureQueryCache: Dict[str, Tuple[float, str]] = {}
_CACHE_TTL_SECONDS = 300


def _getOrCreateFeatureDbConnector(featureDbName: str, userId: str):
    """Reuse a pooled DB connector for the given feature database."""
    if featureDbName in _featureDbConnPool:
        conn = _featureDbConnPool[featureDbName]
        try:
            if conn.connection and not conn.connection.closed:
                return conn
        except Exception as e:
            logger.warning(f"Feature DB connection check failed for {featureDbName}: {e}")
        _featureDbConnPool.pop(featureDbName, None)

    from modules.connectors.connectorDbPostgre import DatabaseConnector
    from modules.shared.configuration import APP_CONFIG
    conn = DatabaseConnector(
        dbHost=APP_CONFIG.get("DB_HOST", "localhost"),
        dbDatabase=featureDbName,
        dbUser=APP_CONFIG.get("DB_USER"),
        dbPassword=APP_CONFIG.get("DB_PASSWORD_SECRET"),
        dbPort=int(APP_CONFIG.get("DB_PORT", 5432)),
        userId=userId,
    )
    _featureDbConnPool[featureDbName] = conn
    return conn


def clearFeatureQueryCache(featureInstanceId: Optional[str] = None) -> int:
    """Clear the feature data query cache. If featureInstanceId given, only for that instance."""
    if featureInstanceId:
        prefix = f"{featureInstanceId}:"
        keys = [k for k in _featureQueryCache if k.startswith(prefix)]
    else:
        keys = list(_featureQueryCache.keys())
    for k in keys:
        del _featureQueryCache[k]
    logger.info(f"Feature query cache cleared: {len(keys)} entries removed (instance={featureInstanceId or 'all'})")
    return len(keys)


def _registerFeatureSubAgentTools(registry: ToolRegistry, services):
    """Auto-extracted from registerCoreTools."""
    # ---- Feature Data Sub-Agent tool ----

    async def _queryFeatureInstance(args: Dict[str, Any], context: Dict[str, Any]):
        """Delegate a question to the Feature Data Sub-Agent."""
        featureInstanceId = args.get("featureInstanceId", "")
        question = args.get("question", "")
        if not featureInstanceId or not question:
            return ToolResult(
                toolCallId="", toolName="queryFeatureInstance",
                success=False, error="featureInstanceId and question are required",
            )
        try:
            from modules.serviceCenter.services.serviceAgent.featureDataAgent import runFeatureDataAgent
            from modules.datamodels.datamodelFeatureDataSource import FeatureDataSource
            from modules.interfaces.interfaceDbApp import getRootInterface

            rootIf = getRootInterface()
            instance = rootIf.getFeatureInstance(featureInstanceId)
            if not instance:
                return ToolResult(
                    toolCallId="", toolName="queryFeatureInstance",
                    success=False, error=f"Feature instance {featureInstanceId} not found",
                )

            featureCode = instance.featureCode
            mandateId = instance.mandateId or ""
            instanceLabel = instance.label or ""
            userId = context.get("userId", "")
            workspaceInstanceId = context.get("featureInstanceId", "")
            requestLang = None
            if userId:
                langUser = rootIf.getUser(userId)
                if langUser:
                    requestLang = getattr(langUser, "language", None)

            rootDbConn = rootIf.db if hasattr(rootIf, "db") else None
            if rootDbConn is None:
                return ToolResult(
                    toolCallId="", toolName="queryFeatureInstance",
                    success=False, error="No database connector available",
                )

            featureDataSources = rootDbConn.getRecordset(
                FeatureDataSource,
                recordFilter={"featureInstanceId": featureInstanceId, "workspaceInstanceId": workspaceInstanceId},
            )

            _anySourceNeutralize = any(
                bool(ds.get("neutralize", False) if isinstance(ds, dict) else getattr(ds, "neutralize", False))
                for ds in (featureDataSources or [])
            )

            from modules.security.rbacCatalog import getCatalogService
            catalog = getCatalogService()
            tableFilters = {}
            if not featureDataSources:
                selectedTables = catalog.getDataObjects(featureCode)
            else:
                allObjs = {o["meta"]["table"]: o for o in catalog.getDataObjects(featureCode) if "meta" in o and "table" in o.get("meta", {})}
                selectedTables = []
                _wildcardExpanded = False
                for ds in featureDataSources:
                    tn = ds.get("tableName", "")
                    ok = ds.get("objectKey", "")
                    if ok.endswith(".*") or (not tn and ok):
                        selectedTables = list(allObjs.values())
                        _wildcardExpanded = True
                        break
                    if tn in allObjs:
                        selectedTables.append(allObjs[tn])
                if not _wildcardExpanded:
                    for ds in featureDataSources:
                        rf = ds.get("recordFilter")
                        if rf and isinstance(rf, dict) and ds.get("tableName"):
                            tableFilters[ds["tableName"]] = rf

            if not selectedTables:
                return ToolResult(
                    toolCallId="", toolName="queryFeatureInstance",
                    success=False, error=f"No data tables available for feature '{featureCode}'",
                )

            cacheKey = f"{featureInstanceId}:{hashlib.md5(question.encode()).hexdigest()}"
            if cacheKey in _featureQueryCache:
                cachedAt, cachedResult = _featureQueryCache[cacheKey]
                if time.time() - cachedAt < _CACHE_TTL_SECONDS:
                    return ToolResult(
                        toolCallId="", toolName="queryFeatureInstance",
                        success=True, data=cachedResult,
                    )

            featureDbName = f"poweron_{featureCode.lower()}"
            featureDbConn = _getOrCreateFeatureDbConnector(featureDbName, userId or "agent")

            aiService = services.ai if hasattr(services, "ai") else None
            if aiService is None:
                return ToolResult(
                    toolCallId="", toolName="queryFeatureInstance",
                    success=False, error="AI service not available for sub-agent",
                )

            async def _subAgentAiCall(req):
                if _anySourceNeutralize:
                    req.requireNeutralization = True
                return await aiService.callAi(req)

            answer = await runFeatureDataAgent(
                question=question,
                featureInstanceId=featureInstanceId,
                featureCode=featureCode,
                selectedTables=selectedTables,
                mandateId=mandateId,
                userId=userId,
                aiCallFn=_subAgentAiCall,
                dbConnector=featureDbConn,
                instanceLabel=instanceLabel,
                tableFilters=tableFilters,
                requestLang=requestLang,
            )

            _featureQueryCache[cacheKey] = (time.time(), answer)

            return ToolResult(
                toolCallId="", toolName="queryFeatureInstance",
                success=True, data=answer,
            )
        except Exception as e:
            logger.error(f"queryFeatureInstance failed: {e}", exc_info=True)
            return ToolResult(
                toolCallId="", toolName="queryFeatureInstance",
                success=False, error=str(e),
            )

    registry.register(
        "queryFeatureInstance", _queryFeatureInstance,
        description=(
            "Query data from a feature instance (e.g. Trustee, CommCoach). "
            "Delegates to a sub-agent that knows the feature schema. "
            "Requires the feature instance id from attached feature data sources. "
            "Ask one precise, self-contained question per call."
        ),
        parameters={
            "type": "object",
            "properties": {
                "featureInstanceId": {"type": "string", "description": "ID of the feature instance to query"},
                "question": {"type": "string", "description": "What data to find or analyze from this feature instance"},
            },
            "required": ["featureInstanceId", "question"]
        },
        readOnly=True
    )
