# Copyright (c) 2025 Patrick Motsch
# All rights reserved.
"""Agent service: entry point for running AI agents with tool use."""

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

from modules.datamodels.datamodelAi import (
    AiCallRequest, AiCallOptions, AiCallResponse, OperationTypeEnum
)
from modules.serviceCenter.services.serviceAgent.datamodelAgent import (
    AgentConfig, AgentEvent, AgentEventTypeEnum
)
from modules.serviceCenter.services.serviceAgent.toolRegistry import ToolRegistry
from modules.serviceCenter.services.serviceAgent.agentLoop import runAgentLoop
from modules.serviceCenter.services.serviceAgent.actionToolAdapter import ActionToolAdapter
from modules.serviceCenter.services.serviceAgent.coreTools import registerCoreTools
from modules.serviceCenter.services.serviceBilling.mainServiceBilling import (
    getService as getBillingService,
    InsufficientBalanceException,
    BillingContextError
)

logger = logging.getLogger(__name__)

class _ServicesAdapter:
    """Adapter providing service access from (context, get_service)."""

    def __init__(self, context, getService: Callable[[str], Any]):
        self._context = context
        self._getService = getService
        self.user = context.user
        self.mandateId = context.mandate_id
        self.featureInstanceId = context.feature_instance_id

    @property
    def workflow(self):
        return self._context.workflow

    @property
    def ai(self):
        return self._getService("ai")

    @property
    def chat(self):
        return self._getService("chat")

    @property
    def streaming(self):
        return self._getService("streaming")

    @property
    def billing(self):
        return self._getService("billing")

    @property
    def utils(self):
        return self._getService("utils")

    @property
    def extraction(self):
        return self._getService("extraction")

    def getService(self, name: str):
        """Access any service by name."""
        return self._getService(name)

    @property
    def featureCode(self) -> Optional[str]:
        w = self.workflow
        if w and hasattr(w, "feature") and w.feature:
            return getattr(w.feature, "code", None)
        return getattr(w, "featureCode", None) if w else None


class AgentService:
    """Service for running AI agents with ReAct loop and tool use.

    Registered as IMPORTABLE_SERVICE with objectKey 'service.agent'.
    Uses serviceAi for model selection/billing, streaming for SSE events.
    """

    def __init__(self, context, get_service: Callable[[str], Any]):
        self._context = context
        self._getService = get_service
        self.services = _ServicesAdapter(context, get_service)

    async def runAgent(
        self,
        prompt: str,
        fileIds: List[str] = None,
        config: AgentConfig = None,
        toolSet: str = "core",
        workflowId: str = None,
        additionalTools: List[Dict[str, Any]] = None,
        userLanguage: str = "",
        conversationHistory: List[Dict[str, Any]] = None,
        buildRagContextFn: Callable = None,
        systemPromptOverride: str = None,
    ) -> AsyncGenerator[AgentEvent, None]:
        """Run an agent with the given prompt and tools.

        Args:
            prompt: User prompt
            fileIds: Optional list of file IDs to include as context
            config: Agent configuration
            toolSet: Which tool set to activate
            workflowId: Workflow ID for tracking and billing
            additionalTools: Extra tool definitions to register dynamically
            userLanguage: ISO 639-1 language code; falls back to user.language from profile
            conversationHistory: Prior messages for follow-up context
            buildRagContextFn: Optional custom RAG context builder (overrides default)
            systemPromptOverride: Optional system prompt override (replaces generated prompt)

        Yields:
            AgentEvent for each step (SSE-ready)
        """
        if config is None:
            config = AgentConfig(toolSet=toolSet)

        if workflowId is None:
            workflowId = getattr(self.services.workflow, "id", "unknown") if self.services.workflow else "unknown"

        resolvedLanguage = userLanguage or getattr(self.services.user, "language", "") or "de"

        enrichedPrompt = await self._enrichPromptWithFiles(prompt, fileIds)

        toolRegistry = self._buildToolRegistry(config)

        aiCallFn = self._createAiCallFn()
        aiCallStreamFn = self._createAiCallStreamFn()
        getWorkflowCostFn = self._createGetWorkflowCostFn(workflowId)
        if buildRagContextFn is None:
            buildRagContextFn = self._createBuildRagContextFn()
        persistRoundMemoryFn = self._createPersistRoundMemoryFn(workflowId)
        getExternalMemoryKeysFn = self._createGetExternalMemoryKeysFn(workflowId)

        async for event in runAgentLoop(
            prompt=enrichedPrompt,
            toolRegistry=toolRegistry,
            config=config,
            aiCallFn=aiCallFn,
            getWorkflowCostFn=getWorkflowCostFn,
            workflowId=workflowId,
            userId=self.services.user.id if self.services.user else "",
            featureInstanceId=self.services.featureInstanceId or "",
            buildRagContextFn=buildRagContextFn,
            mandateId=self.services.mandateId or "",
            aiCallStreamFn=aiCallStreamFn,
            userLanguage=resolvedLanguage,
            conversationHistory=conversationHistory,
            persistRoundMemoryFn=persistRoundMemoryFn,
            getExternalMemoryKeysFn=getExternalMemoryKeysFn,
            systemPromptOverride=systemPromptOverride,
        ):
            if event.type == AgentEventTypeEnum.AGENT_SUMMARY:
                await self._persistTrace(workflowId, event.data or {})
            if event.type != AgentEventTypeEnum.CHUNK:
                logger.debug(f"runAgent yielding event type={event.type}")
            yield event
        logger.info(f"runAgent loop completed for workflow {workflowId}")

    async def _enrichPromptWithFiles(self, prompt: str, fileIds: List[str] = None) -> str:
        """Resolve file metadata + FileContentIndex for attached fileIds and prepend to prompt.

        The FileContentIndex is produced by the upload pipeline (AI-free extraction)
        and tells the agent exactly which content objects (text, images, tables, etc.)
        exist inside a file, so the agent can work with them directly via tools.
        """
        if not fileIds:
            return prompt
        try:
            chatService = self.services.chat
            knowledgeDb = None
            try:
                from modules.interfaces.interfaceDbKnowledge import getInterface as _getKnowledgeInterface
                knowledgeDb = _getKnowledgeInterface()
            except Exception:
                pass

            fileDescriptions = []
            for fid in fileIds:
                try:
                    info = chatService.getFileInfo(fid)

                    if not info:
                        folderInfo = chatService.interfaceDbComponent.getFolder(fid)
                        if folderInfo:
                            folderName = folderInfo.get("name", fid)
                            folderFiles = chatService.listFiles(folderId=fid)
                            desc = f"### Folder: {folderName}\n  - id: {fid}\n  - type: folder\n  - contains: {len(folderFiles)} file(s)"
                            if folderFiles:
                                desc += "\n  - files:"
                                for ff in folderFiles[:30]:
                                    ffName = ff.get("fileName", "?")
                                    ffId = ff.get("id", "?")
                                    ffMime = ff.get("mimeType", "?")
                                    ffSize = ff.get("fileSize", ff.get("size", "?"))
                                    desc += f"\n    * {ffName} (id: {ffId}, type: {ffMime}, size: {ffSize} bytes)"
                                if len(folderFiles) > 30:
                                    desc += f"\n    ... and {len(folderFiles) - 30} more files"
                            desc += f'\nUse `listFiles(folderId="{fid}")` to get the full file list, then `readFile(fileId)` to read individual files.'
                            fileDescriptions.append(desc)
                            continue
                        fileDescriptions.append(f"### File id: {fid}")
                        continue

                    fileName = info.get("fileName", fid)
                    mimeType = info.get("mimeType", "unknown")
                    fileSize = info.get("size", "?")

                    desc = f"### File: {fileName}\n  - id: {fid}\n  - type: {mimeType}\n  - size: {fileSize} bytes"

                    if knowledgeDb:
                        contentIndex = knowledgeDb.getFileContentIndex(fid)
                        if contentIndex:
                            structure = contentIndex.get("structure", {})
                            totalObjects = contentIndex.get("totalObjects", 0)
                            desc += f"\n  - indexed: yes ({totalObjects} content objects)"
                            if structure:
                                structParts = []
                                for key, val in structure.items():
                                    if isinstance(val, (int, str)):
                                        structParts.append(f"{key}: {val}")
                                if structParts:
                                    desc += f"\n  - structure: {', '.join(structParts)}"

                            objectSummary = contentIndex.get("objectSummary", [])
                            if objectSummary:
                                desc += "\n  - content objects:"
                                for obj in objectSummary[:20]:
                                    objType = obj.get("type", obj.get("contentType", "?"))
                                    objRef = obj.get("ref", {})
                                    objLabel = objRef.get("location", "") if isinstance(objRef, dict) else ""
                                    objId = obj.get("id", obj.get("contentObjectId", ""))
                                    desc += f"\n    * [{objType}] {objLabel}" + (f" (id: {objId})" if objId else "")
                                if len(objectSummary) > 20:
                                    desc += f"\n    ... and {len(objectSummary) - 20} more objects"
                        else:
                            desc += "\n  - indexed: no (use readFile to trigger extraction)"

                    fileDescriptions.append(desc)
                except Exception:
                    fileDescriptions.append(f"### File id: {fid}")

            if fileDescriptions:
                header = (
                    "## Attached Files & Folders\n"
                    "These files/folders have been uploaded and processed through the extraction pipeline.\n"
                    "Use `readFile(fileId)` to read text content, `readContentObjects(fileId)` for structured access, "
                    "or `describeImage(fileId)` for image analysis.\n"
                    "For folders, use `listFiles(folderId)` to get the files inside, then `readFile(fileId)` for each.\n"
                    "For large PDFs/DOCX, avoid huge `renderDocument` tool JSON: build markdown with "
                    "`writeFile` (create + append), then `renderDocument(sourceFileId=that file id, outputFormat=...)`.\n"
                    "For small docs you may pass `content` inline. Embed images with `![alt](file:fileId)` in markdown.\n\n"
                )
                header += "\n\n".join(fileDescriptions)
                return f"{header}\n\n---\n\nUser request: {prompt}"
        except Exception as e:
            logger.warning(f"Could not enrich prompt with file metadata: {e}")
        return prompt

    def _buildToolRegistry(self, config: AgentConfig) -> ToolRegistry:
        """Build a tool registry with core tools, ActionToolAdapter tools, and toolbox-activated tools."""
        registry = ToolRegistry()

        registerCoreTools(registry, self.services)

        try:
            from modules.workflows.processing.core.actionExecutor import ActionExecutor
            actionExecutor = ActionExecutor(self.services)
            adapter = ActionToolAdapter(actionExecutor)
            adapter.registerAll(registry)
        except Exception as e:
            logger.warning(f"Could not register action tools: {e}")

        self._activateToolboxes(registry, config)

        return registry

    def _activateToolboxes(self, registry: ToolRegistry, config: AgentConfig) -> None:
        """Activate toolboxes dynamically based on user connections and config."""
        try:
            from modules.serviceCenter.services.serviceAgent.toolboxRegistry import getToolboxRegistry
            tbRegistry = getToolboxRegistry()

            userConnections: List[str] = []
            try:
                connectionService = self._getService("connection")
                if connectionService and hasattr(connectionService, "getConnections"):
                    connections = connectionService.getConnections() or []
                    userConnections = [c.get("authority", "") for c in connections if c.get("authority")]
            except Exception as e:
                logger.debug("Could not resolve user connections for toolbox activation: %s", e)

            activeToolboxes = tbRegistry.getActiveToolboxes(userConnections)
            activatedIds = [tb.id for tb in activeToolboxes]
            logger.info("Toolbox activation: connections=%s -> active toolboxes=%s", userConnections, activatedIds)

            for tb in activeToolboxes:
                if tb.id == "workflow":
                    try:
                        from modules.serviceCenter.services.serviceAgent.workflowTools import getWorkflowToolDefinitions
                        for toolDef in getWorkflowToolDefinitions():
                            registry.register(toolDef)
                        logger.info("Registered %d workflow tools from toolbox", len(getWorkflowToolDefinitions()))
                    except Exception as e:
                        logger.warning("Could not register workflow tools: %s", e)
        except Exception as e:
            logger.warning("Toolbox activation failed: %s", e)

    async def _persistTrace(self, workflowId: str, summaryData: Dict[str, Any]):
        """Persist the agent trace and workflow artifacts in the knowledge store."""
        try:
            knowledgeService = self._getService("knowledge")
            userId = self.services.user.id if self.services.user else ""
            featureInstanceId = self.services.featureInstanceId or ""

            import json
            traceValue = json.dumps(summaryData, default=str)

            await knowledgeService.storeEntity(
                workflowId=workflowId,
                userId=userId,
                featureInstanceId=featureInstanceId,
                key="_agentTrace",
                value=traceValue,
                source="agent",
            )

            artifacts = summaryData.get("artifacts", "")
            if artifacts:
                await knowledgeService.storeEntity(
                    workflowId=workflowId,
                    userId=userId,
                    featureInstanceId=featureInstanceId,
                    key="_workflowArtifacts",
                    value=artifacts,
                    source="agent",
                )
                logger.info(f"Persisted workflow artifacts for workflow {workflowId}")

            logger.info(f"Persisted agent trace for workflow {workflowId}")
        except Exception as e:
            logger.warning(f"Could not persist agent trace: {e}")

    def _createAiCallFn(self) -> Callable[[AiCallRequest], AiCallResponse]:
        """Create the AI call function that wraps serviceAi with billing."""
        ctxNeutralization = getattr(self._context, "requireNeutralization", None)
        async def _aiCallFn(request: AiCallRequest) -> AiCallResponse:
            if ctxNeutralization is not None and request.requireNeutralization is None:
                request.requireNeutralization = ctxNeutralization
            aiService = self.services.ai
            return await aiService.callAi(request)
        return _aiCallFn

    def _createAiCallStreamFn(self):
        """Create the streaming AI call function. Yields str deltas, then AiCallResponse."""
        ctxNeutralization = getattr(self._context, "requireNeutralization", None)
        async def _aiCallStreamFn(request: AiCallRequest):
            if ctxNeutralization is not None and request.requireNeutralization is None:
                request.requireNeutralization = ctxNeutralization
            aiService = self.services.ai
            async for chunk in aiService.callAiStream(request):
                yield chunk
        return _aiCallStreamFn

    def _createGetWorkflowCostFn(self, workflowId: str) -> Callable[[], float]:
        """Create a function that returns the current workflow cost."""
        async def _getWorkflowCost() -> float:
            try:
                billingService = self.services.billing
                return await billingService.getWorkflowCost(workflowId)
            except Exception:
                return 0.0
        return _getWorkflowCost

    def _createBuildRagContextFn(self):
        """Create the RAG context builder function that delegates to KnowledgeService."""
        async def _buildRagContext(
            currentPrompt: str, workflowId: str, userId: str,
            featureInstanceId: str, mandateId: str, **kwargs
        ) -> str:
            try:
                knowledgeService = self.services.getService("knowledge")
                workflowHintItems = _buildWorkflowHintItems(
                    self.services, workflowId
                )
                return await knowledgeService.buildAgentContext(
                    currentPrompt=currentPrompt,
                    workflowId=workflowId,
                    userId=userId,
                    featureInstanceId=featureInstanceId,
                    mandateId=mandateId,
                    workflowHintItems=workflowHintItems,
                    isSysAdmin=getattr(self.services.user, "isSysAdmin", False),
                )
            except Exception as e:
                logger.debug(f"RAG context not available: {e}")
                return ""
        return _buildRagContext

    def _createPersistRoundMemoryFn(self, workflowId: str):
        """Create callback that persists RoundMemory entries after tool execution."""
        from modules.serviceCenter.services.serviceAgent.agentLoop import _classifyToolResult
        from modules.datamodels.datamodelKnowledge import RoundMemory

        async def _persistRoundMemory(
            toolCalls, results, textContent: str, roundNumber: int
        ):
            try:
                knowledgeService = self.services.getService("knowledge")
            except Exception:
                return
            knowledgeDb = knowledgeService._knowledgeDb

            for tc, result in zip(toolCalls, results):
                if not result.success:
                    continue
                classified = _classifyToolResult(tc, result)
                if not classified:
                    continue

                summary = classified["summary"]
                embedding = await knowledgeService._embedSingle(summary[:500]) if summary else []

                mem = RoundMemory(
                    workflowId=workflowId,
                    roundNumber=roundNumber,
                    memoryType=classified["memoryType"],
                    key=classified["key"],
                    summary=summary,
                    fullData=classified.get("fullData"),
                    fileIds=classified.get("fileIds", []),
                    embedding=embedding if embedding else None,
                )
                knowledgeDb.storeRoundMemory(mem)

        return _persistRoundMemory

    def _createGetExternalMemoryKeysFn(self, workflowId: str):
        """Create callback that returns RoundMemory keys for summarization hints."""
        def _getKeys() -> List[str]:
            try:
                knowledgeService = self.services.getService("knowledge")
                memories = knowledgeService._knowledgeDb.getRoundMemories(workflowId)
                return [m.get("key", "") for m in memories if m.get("key")]
            except Exception:
                return []
        return _getKeys


def _buildWorkflowHintItems(
    services, currentWorkflowId: str
) -> List[Dict[str, Any]]:
    """Build a compact list of other workflows for the RAG cross-workflow hint.

    Returns key-value items like:
        key="Pendenzenliste Excel (3 msgs)"  value="last: 2h ago"
    Limited to 10 most recent other workflows to keep the hint small.
    """
    try:
        chatInterface = services.chat.interfaceDbChat
        allWorkflows = chatInterface.getWorkflows() or []
    except Exception:
        return []

    others = [w for w in allWorkflows if w.get("id") != currentWorkflowId]
    if not others:
        return []

    import time as _time
    now = _time.time()
    others.sort(key=lambda w: w.get("sysCreatedAt") or w.get("startedAt") or 0, reverse=True)
    others = others[:10]

    items = []
    for wf in others:
        name = wf.get("name") or "(unnamed)"
        createdAt = wf.get("sysCreatedAt") or wf.get("startedAt") or 0
        ageSec = now - createdAt if createdAt else 0
        if ageSec < 3600:
            ageStr = f"{int(ageSec / 60)}m ago"
        elif ageSec < 86400:
            ageStr = f"{int(ageSec / 3600)}h ago"
        else:
            ageStr = f"{int(ageSec / 86400)}d ago"

        wfId = wf.get("id", "")
        items.append({
            "key": f"{name} (id: {wfId})",
            "value": ageStr,
        })

    countLabel = f"{len(allWorkflows) - 1} other conversation(s)"
    if len(allWorkflows) - 1 > 10:
        countLabel += f" (showing 10 newest)"
    items.insert(0, {"key": countLabel, "value": "use listWorkflowHistory to browse"})
    return items