# Copyright (c) 2025 Patrick Motsch
# All rights reserved.
"""DataSource convenience tools (browse, search, download from external sources)."""

import logging
from typing import Any, Dict, List, Optional

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__)


def _registerDataSourceTools(registry: ToolRegistry, services):
    """Auto-extracted from registerCoreTools."""

    def _buildResolverDb():
        """Build a DB adapter that ConnectorResolver can use to load UserConnections.
        interfaceDbApp has getUserConnectionById; ConnectorResolver expects getUserConnection."""
        chatService = services.chat
        appIf = getattr(chatService, "interfaceDbApp", None)
        if appIf and hasattr(appIf, "getUserConnectionById"):
            class _Adapter:
                def __init__(self, app):
                    self._app = app
                def getUserConnection(self, connectionId: str):
                    return self._app.getUserConnectionById(connectionId)
            return _Adapter(appIf)
        return getattr(chatService, "interfaceDbComponent", None)

    # ---- DataSource convenience tools ----
    _SOURCE_TYPE_TO_SERVICE = {
        "sharepointFolder": "sharepoint",
        "onedriveFolder": "onedrive",
        "outlookFolder": "outlook",
        "googleDriveFolder": "drive",
        "gmailFolder": "gmail",
        "ftpFolder": "files",
        "clickupList": "clickup",
    }

    async def _resolveDataSource(dsId: str):
        """Resolve a DataSource record and return (connectionId, service, path, neutralize) or raise."""
        chatService = services.chat
        ds = chatService.getDataSource(dsId) if hasattr(chatService, "getDataSource") else None
        if not ds:
            raise ValueError(f"DataSource '{dsId}' not found")
        connectionId = ds.get("connectionId", "")
        sourceType = ds.get("sourceType", "")
        path = ds.get("path", "/")
        label = ds.get("label", "")
        neutralize = bool(ds.get("neutralize", False))
        service = _SOURCE_TYPE_TO_SERVICE.get(sourceType, sourceType)
        if not connectionId:
            raise ValueError(f"DataSource '{dsId}' has no connectionId")
        logger.info(f"Resolved DataSource '{dsId}' ({label}): sourceType={sourceType}, service={service}, connectionId={connectionId}, path={path[:80]}, neutralize={neutralize}")
        return connectionId, service, path, neutralize

    _MAIL_SERVICES = {"outlook", "gmail"}

    async def _browseDataSource(args: Dict[str, Any], context: Dict[str, Any]):
        dsId = args.get("dataSourceId", "")
        subPath = args.get("subPath", "")
        directConnId = args.get("connectionId", "")
        directService = args.get("service", "")
        if not dsId and not (directConnId and directService):
            return ToolResult(toolCallId="", toolName="browseDataSource", success=False,
                              error="Provide either dataSourceId OR connectionId+service")
        try:
            if dsId:
                connectionId, service, basePath, _neutralize = await _resolveDataSource(dsId)
            else:
                connectionId, service, basePath = directConnId, directService, args.get("path", "/")
            if subPath:
                if subPath.startswith("/"):
                    browsePath = subPath
                else:
                    browsePath = f"{basePath.rstrip('/')}/{subPath}"
            else:
                browsePath = basePath
            from modules.connectors.connectorResolver import ConnectorResolver
            resolver = ConnectorResolver(
                services.getService("security"),
                _buildResolverDb(),
            )
            adapter = await resolver.resolveService(connectionId, service)
            entries = await adapter.browse(browsePath, filter=args.get("filter"))
            if not entries:
                return ToolResult(toolCallId="", toolName="browseDataSource", success=True, data="Empty directory.")
            lines = []
            for e in entries:
                prefix = "[DIR]" if e.isFolder else "[FILE]"
                sizeInfo = f" ({e.size} bytes)" if e.size else ""
                lines.append(f"- {prefix} {e.name}{sizeInfo}  path: {e.path}")
            result = "\n".join(lines)
            if service in _MAIL_SERVICES:
                result += "\n\nIMPORTANT: These are email subjects only. To read the full email content, use downloadFromDataSource with the path, then readFile on the returned file ID."
            return ToolResult(toolCallId="", toolName="browseDataSource", success=True, data=result)
        except Exception as e:
            return ToolResult(toolCallId="", toolName="browseDataSource", success=False, error=str(e))

    async def _searchDataSource(args: Dict[str, Any], context: Dict[str, Any]):
        dsId = args.get("dataSourceId", "")
        directConnId = args.get("connectionId", "")
        directService = args.get("service", "")
        query = args.get("query", "")
        if not query:
            return ToolResult(toolCallId="", toolName="searchDataSource", success=False, error="query is required")
        if not dsId and not (directConnId and directService):
            return ToolResult(toolCallId="", toolName="searchDataSource", success=False,
                              error="Provide either dataSourceId OR connectionId+service")
        try:
            if dsId:
                connectionId, service, basePath, _neutralize = await _resolveDataSource(dsId)
            else:
                connectionId, service, basePath = directConnId, directService, args.get("path", "/")
            from modules.connectors.connectorResolver import ConnectorResolver
            resolver = ConnectorResolver(
                services.getService("security"),
                _buildResolverDb(),
            )
            adapter = await resolver.resolveService(connectionId, service)
            entries = await adapter.search(query, path=basePath)
            if not entries:
                return ToolResult(toolCallId="", toolName="searchDataSource", success=True, data="No results found.")
            lines = [f"- {e.name} (path: {e.path})" for e in entries]
            result = "\n".join(lines)
            if service in _MAIL_SERVICES:
                result += "\n\nIMPORTANT: These are email subjects only. To read the full email content, use downloadFromDataSource with the path, then readFile on the returned file ID."
            return ToolResult(toolCallId="", toolName="searchDataSource", success=True, data=result)
        except Exception as e:
            return ToolResult(toolCallId="", toolName="searchDataSource", success=False, error=str(e))

    async def _downloadFromDataSource(args: Dict[str, Any], context: Dict[str, Any]):
        dsId = args.get("dataSourceId", "")
        directConnId = args.get("connectionId", "")
        directService = args.get("service", "")
        filePath = args.get("filePath", "")
        fileName = args.get("fileName", "")
        if not filePath:
            return ToolResult(toolCallId="", toolName="downloadFromDataSource", success=False, error="filePath is required")
        if not dsId and not (directConnId and directService):
            return ToolResult(toolCallId="", toolName="downloadFromDataSource", success=False,
                              error="Provide either dataSourceId OR connectionId+service")
        try:
            from modules.connectors.connectorResolver import ConnectorResolver
            from modules.connectors.connectorProviderBase import DownloadResult as _DR
            _sourceNeutralize = False
            if dsId:
                connectionId, service, basePath, _sourceNeutralize = await _resolveDataSource(dsId)
            else:
                connectionId, service, basePath = directConnId, directService, "/"
            fullPath = filePath if filePath.startswith("/") else f"{basePath.rstrip('/')}/{filePath}"
            resolver = ConnectorResolver(
                services.getService("security"),
                _buildResolverDb(),
            )
            adapter = await resolver.resolveService(connectionId, service)
            result = await adapter.download(fullPath)

            if isinstance(result, _DR):
                fileBytes = result.data
                fileName = result.fileName or fileName
            else:
                fileBytes = result

            if not fileBytes:
                return ToolResult(toolCallId="", toolName="downloadFromDataSource", success=False, error="Download returned empty")

            if not fileName or "." not in fileName:
                pathSegment = fullPath.split("/")[-1] or "downloaded_file"
                fileName = fileName or pathSegment
                if "." not in fileName:
                    try:
                        entries = await adapter.browse(basePath)
                        for entry in entries:
                            if getattr(entry, "path", "") == filePath or getattr(entry, "path", "").endswith(filePath):
                                if "." in entry.name:
                                    fileName = entry.name
                                    break
                    except Exception as e:
                        logger.warning(f"downloadFromDataSource: browse for filename failed: {e}")
                if "." not in fileName:
                    if fileBytes[:4] == b"%PDF":
                        fileName = f"{fileName}.pdf"
                    elif fileBytes[:2] == b"PK":
                        fileName = f"{fileName}.zip"
            chatService = services.chat
            fileItem, _ = chatService.interfaceDbComponent.saveUploadedFile(fileBytes, fileName)
            fiId = context.get("featureInstanceId") or (services.featureInstanceId if services else "")
            if fiId:
                chatService.interfaceDbComponent.updateFile(fileItem.id, {"featureInstanceId": fiId})
            if _sourceNeutralize:
                chatService.interfaceDbComponent.updateFile(fileItem.id, {"neutralize": True})
            tempFolderId = _getOrCreateTempFolder(chatService)
            if tempFolderId:
                chatService.interfaceDbComponent.updateFile(fileItem.id, {"folderId": tempFolderId})
            ext = fileName.rsplit(".", 1)[-1].lower() if "." in fileName else ""
            hint = "Use readFile to read the text content." if ext in ("doc", "docx", "txt", "csv", "json", "xml", "html", "md", "rtf", "odt", "xls", "xlsx", "pptx", "pdf", "eml", "msg") else "Use readFile to access the content."
            return ToolResult(
                toolCallId="", toolName="downloadFromDataSource", success=True,
                data=f"Downloaded '{fileName}' ({len(fileBytes)} bytes) → local file id: {fileItem.id}. {hint}"
            )
        except Exception as e:
            return ToolResult(toolCallId="", toolName="downloadFromDataSource", success=False, error=str(e))

    registry.register(
        "browseDataSource", _browseDataSource,
        description=(
            "Browse files and folders in a data source. Accepts either:\n"
            "- dataSourceId (for attached data sources shown in the prompt), OR\n"
            "- connectionId + service (for direct connection access via listConnections)."
        ),
        parameters={
            "type": "object",
            "properties": {
                "dataSourceId": {"type": "string", "description": "DataSource ID (from attached data sources)"},
                "connectionId": {"type": "string", "description": "UserConnection ID (alternative to dataSourceId)"},
                "service": {"type": "string", "description": "Service name (alternative to dataSourceId, e.g. sharepoint, onedrive)"},
                "path": {"type": "string", "description": "Root path (used with connectionId+service)"},
                "subPath": {"type": "string", "description": "Sub-path within the data source to browse"},
                "filter": {"type": "string", "description": "Filter pattern (e.g. '*.pdf')"},
            },
        },
        readOnly=True,
    )

    registry.register(
        "searchDataSource", _searchDataSource,
        description=(
            "Search for files within a data source. Accepts either dataSourceId OR connectionId+service."
        ),
        parameters={
            "type": "object",
            "properties": {
                "dataSourceId": {"type": "string", "description": "DataSource ID"},
                "connectionId": {"type": "string", "description": "UserConnection ID (alternative to dataSourceId)"},
                "service": {"type": "string", "description": "Service name (alternative to dataSourceId)"},
                "path": {"type": "string", "description": "Scope path (used with connectionId+service)"},
                "query": {"type": "string", "description": "Search query"},
            },
            "required": ["query"],
        },
        readOnly=True,
    )

    registry.register(
        "downloadFromDataSource", _downloadFromDataSource,
        description=(
            "Download a file or email from a data source into local storage. Returns a local file ID "
            "to read with readFile. Accepts either dataSourceId OR connectionId+service. "
            "For email sources (Outlook, Gmail), browse/search only return subjects -- use this to get full content."
        ),
        parameters={
            "type": "object",
            "properties": {
                "dataSourceId": {"type": "string", "description": "DataSource ID"},
                "connectionId": {"type": "string", "description": "UserConnection ID (alternative to dataSourceId)"},
                "service": {"type": "string", "description": "Service name (alternative to dataSourceId)"},
                "filePath": {"type": "string", "description": "Path of the file to download (from browseDataSource results)"},
                "fileName": {"type": "string", "description": "File name with extension (e.g. 'report.pdf')"},
            },
            "required": ["dataSourceId", "filePath"],
        },
        readOnly=False,
    )
