# Copyright (c) 2025 Patrick Motsch
# All rights reserved.
"""Document and vision tools (containers, content objects, image description)."""

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 _registerDocumentTools(registry: ToolRegistry, services):
    """Auto-extracted from registerCoreTools."""
    # ---- Document tools (Smart Documents / Container Handling) ----

    async def _browseContainer(args: Dict[str, Any], context: Dict[str, Any]):
        fileId = args.get("fileId", "")
        if not fileId:
            return ToolResult(toolCallId="", toolName="browseContainer", success=False, error="fileId is required")
        try:
            knowledgeService = services.getService("knowledge")
            index = knowledgeService.getFileContentIndex(fileId)
            if not index:
                return ToolResult(toolCallId="", toolName="browseContainer", success=True, data="No content index available for this file. It may not have been indexed yet.")
            structure = index.get("structure", {}) if isinstance(index, dict) else {}
            objectSummary = index.get("objectSummary", []) if isinstance(index, dict) else []
            totalObjects = index.get("totalObjects", 0) if isinstance(index, dict) else 0

            result = f"File: {index.get('fileName', '?')} ({index.get('mimeType', '?')})\n"
            result += f"Total content objects: {totalObjects}\n"

            sections = structure.get("sections", [])
            if sections:
                result += "\nSections:\n"
                for s in sections:
                    result += f"  [{s.get('id', '?')}] {s.get('title', 'Untitled')} (pages {s.get('startPage', '?')}-{s.get('endPage', '?')})\n"

            if structure.get("pageMap"):
                pages = len(structure["pageMap"])
                result += f"\nPages: {pages}\n"
                imgCount = structure.get("imageCount", 0)
                tableCount = structure.get("tableCount", 0)
                if imgCount:
                    result += f"Images: {imgCount}\n"
                if tableCount:
                    result += f"Tables: {tableCount}\n"

            if structure.get("sheetMap"):
                result += "\nSheets:\n"
                for s in structure["sheetMap"]:
                    result += f"  {s.get('sheetName', '?')} ({s.get('rows', '?')} rows x {s.get('columns', '?')} cols)\n"

            if structure.get("slideMap"):
                result += "\nSlides:\n"
                for s in structure["slideMap"]:
                    result += f"  Slide {s.get('slideIndex', 0) + 1}: {s.get('title', '(no title)')}\n"

            return ToolResult(toolCallId="", toolName="browseContainer", success=True, data=result)
        except Exception as e:
            return ToolResult(toolCallId="", toolName="browseContainer", success=False, error=str(e))

    async def _readContentObjects(args: Dict[str, Any], context: Dict[str, Any]):
        fileId = args.get("fileId", "")
        if not fileId:
            return ToolResult(toolCallId="", toolName="readContentObjects", success=False, error="fileId is required")
        try:
            knowledgeService = services.getService("knowledge")
            filterDict = {}
            if args.get("pageIndex") is not None:
                filterDict["pageIndex"] = args["pageIndex"]
            if args.get("contentType"):
                filterDict["contentType"] = args["contentType"]
            if args.get("sectionId"):
                filterDict["sectionId"] = args["sectionId"]

            objects = await knowledgeService.readContentObjects(fileId, filterDict)
            if not objects:
                return ToolResult(toolCallId="", toolName="readContentObjects", success=True, data="No content objects found with the given filter.")

            result = f"Found {len(objects)} content objects:\n\n"
            for obj in objects[:20]:
                data = obj.get("data", "")
                cType = obj.get("contentType", "?")
                ref = obj.get("contextRef", {})
                location = ref.get("location", "") if isinstance(ref, dict) else ""
                preview = data[:300] if cType == "text" else f"[{cType} data, {len(data)} chars]"
                result += f"[{cType}] {location}: {preview}\n\n"

            if len(objects) > 20:
                result += f"... and {len(objects) - 20} more objects"

            return ToolResult(toolCallId="", toolName="readContentObjects", success=True, data=result)
        except Exception as e:
            return ToolResult(toolCallId="", toolName="readContentObjects", success=False, error=str(e))

    async def _extractContainerItem(args: Dict[str, Any], context: Dict[str, Any]):
        fileId = args.get("fileId", "")
        containerPath = args.get("containerPath", "")
        if not fileId or not containerPath:
            return ToolResult(toolCallId="", toolName="extractContainerItem", success=False, error="fileId and containerPath are required")
        try:
            knowledgeService = services.getService("knowledge")
            result = await knowledgeService.extractContainerItem(fileId, containerPath)
            if result:
                return ToolResult(toolCallId="", toolName="extractContainerItem", success=True, data=str(result))
            return ToolResult(toolCallId="", toolName="extractContainerItem", success=False, error=f"Item '{containerPath}' not found in container index for file {fileId}. On-demand extraction is not yet implemented.")
        except Exception as e:
            return ToolResult(toolCallId="", toolName="extractContainerItem", success=False, error=str(e))

    async def _summarizeContent(args: Dict[str, Any], context: Dict[str, Any]):
        fileId = args.get("fileId", "")
        if not fileId:
            return ToolResult(toolCallId="", toolName="summarizeContent", success=False, error="fileId is required")
        try:
            knowledgeService = services.getService("knowledge")
            filterDict = {}
            if args.get("sectionId"):
                filterDict["sectionId"] = args["sectionId"]
            if args.get("pageIndex") is not None:
                filterDict["pageIndex"] = args["pageIndex"]
            if args.get("contentType"):
                filterDict["contentType"] = args["contentType"]

            objects = await knowledgeService.readContentObjects(fileId, filterDict)
            if not objects:
                return ToolResult(toolCallId="", toolName="summarizeContent", success=True, data="No content found to summarize.")

            textParts = [obj.get("data", "") for obj in objects if obj.get("contentType") != "image"]
            combinedText = "\n\n".join(textParts)[:6000]

            aiService = services.ai
            from modules.datamodels.datamodelAi import AiCallRequest, AiCallOptions, OperationTypeEnum
            summaryRequest = AiCallRequest(
                prompt=f"Summarize the following content concisely:\n\n{combinedText}",
                options=AiCallOptions(operationType=OperationTypeEnum.DATA_ANALYSE),
            )
            response = await aiService.callAi(summaryRequest)
            return ToolResult(toolCallId="", toolName="summarizeContent", success=True, data=response.content)
        except Exception as e:
            return ToolResult(toolCallId="", toolName="summarizeContent", success=False, error=str(e))

    registry.register(
        "browseContainer", _browseContainer,
        description="Browse the structural index of a document (pages, sections, sheets, slides). Use before readContentObjects for targeted reading.",
        parameters={
            "type": "object",
            "properties": {"fileId": {"type": "string", "description": "The file ID to browse"}},
            "required": ["fileId"],
        },
        readOnly=True,
    )

    registry.register(
        "readContentObjects", _readContentObjects,
        description="Read extracted content objects from a file, optionally filtered by page, section, or type. Use browseContainer first to see the structure.",
        parameters={
            "type": "object",
            "properties": {
                "fileId": {"type": "string", "description": "The file ID"},
                "pageIndex": {"type": "integer", "description": "Filter by page index"},
                "sectionId": {"type": "string", "description": "Filter by section ID"},
                "contentType": {"type": "string", "description": "Filter by content type (text, image, etc.)"},
            },
            "required": ["fileId"],
        },
        readOnly=True,
    )

    registry.register(
        "extractContainerItem", _extractContainerItem,
        description="Extract a specific item from a container file (ZIP, nested file). Use browseContainer to see available items.",
        parameters={
            "type": "object",
            "properties": {
                "fileId": {"type": "string", "description": "The container file ID"},
                "containerPath": {"type": "string", "description": "Path within the container"},
            },
            "required": ["fileId", "containerPath"],
        },
        readOnly=True,
    )

    registry.register(
        "summarizeContent", _summarizeContent,
        description="Generate an AI-powered summary of a file's content. Optionally filter by section, page, or content type.",
        parameters={
            "type": "object",
            "properties": {
                "fileId": {"type": "string", "description": "The file ID"},
                "sectionId": {"type": "string", "description": "Optional: summarize only this section"},
                "pageIndex": {"type": "integer", "description": "Optional: summarize only this page"},
                "contentType": {"type": "string", "description": "Optional: filter by content type"},
            },
            "required": ["fileId"],
        },
        readOnly=True,
    )

    # ---- Vision tool ----

    async def _describeImage(args: Dict[str, Any], context: Dict[str, Any]):
        """Analyse an image using AI vision. Uses Knowledge Store chunks produced by Extractors."""
        fileId = args.get("fileId", "")
        prompt = args.get("prompt", "Describe this image in detail. Extract all visible text, tables, and data.")
        pageIndex = args.get("pageIndex")

        if not fileId:
            return ToolResult(toolCallId="", toolName="describeImage", success=False, error="fileId is required")

        try:
            import base64 as _b64

            imageData = None
            mimeType = "image/png"

            knowledgeService = services.getService("knowledge") if hasattr(services, "getService") else None

            # 1) Knowledge Store: image chunks already produced by PdfExtractor / ImageExtractor
            if knowledgeService:
                chunks = knowledgeService._knowledgeDb.getContentChunks(fileId)
                imageChunks = [c for c in (chunks or []) if c.get("contentType") == "image"]
                if pageIndex is not None:
                    imageChunks = [c for c in imageChunks if c.get("contextRef", {}).get("pageIndex") == pageIndex]
                if imageChunks:
                    imageData = imageChunks[0].get("data", "")
                    chunkMime = imageChunks[0].get("contextRef", {}).get("mimeType")
                    if chunkMime:
                        mimeType = chunkMime

            # 2) File not yet indexed -> trigger extraction via ExtractionService, then retry
            if not imageData and knowledgeService and not knowledgeService.isFileIndexed(fileId):
                try:
                    chatService = services.chat
                    fileInfo = chatService.getFileInfo(fileId)
                    fileContent = chatService.getFileContent(fileId)
                    if fileContent and fileInfo:
                        rawData = fileContent.get("data", "")
                        if isinstance(rawData, str) and len(rawData) > 100:
                            rawBytes = _b64.b64decode(rawData)
                        elif isinstance(rawData, bytes):
                            rawBytes = rawData
                        else:
                            rawBytes = None

                        if rawBytes:
                            from modules.serviceCenter.services.serviceExtraction.subRegistry import ExtractorRegistry
                            from modules.serviceCenter.services.serviceExtraction.subPipeline import runExtraction
                            from modules.datamodels.datamodelExtraction import ExtractionOptions

                            fileMime = fileInfo.get("mimeType", "application/octet-stream")
                            fileName = fileInfo.get("fileName", fileId)
                            extracted = runExtraction(
                                ExtractorRegistry(), None,
                                rawBytes, fileName, fileMime, ExtractionOptions(),
                            )

                            contentObjects = []
                            for part in extracted.parts:
                                tg = (part.typeGroup or "").lower()
                                ct = "image" if tg == "image" else "text"
                                if not part.data or not part.data.strip():
                                    continue
                                contentObjects.append({
                                    "contentObjectId": part.id,
                                    "contentType": ct,
                                    "data": part.data,
                                    "contextRef": {"containerPath": fileName, "location": part.label, **(part.metadata or {})},
                                })

                            if contentObjects:
                                _diFiId, _diMId = _resolveFileScope(fileId, context)
                                await knowledgeService.indexFile(
                                    fileId=fileId, fileName=fileName, mimeType=fileMime,
                                    userId=context.get("userId", ""), contentObjects=contentObjects,
                                    featureInstanceId=_diFiId,
                                    mandateId=_diMId,
                                )

                            chunks = knowledgeService._knowledgeDb.getContentChunks(fileId)
                            imageChunks = [c for c in (chunks or []) if c.get("contentType") == "image"]
                            if pageIndex is not None:
                                imageChunks = [c for c in imageChunks if c.get("contextRef", {}).get("pageIndex") == pageIndex]
                            if imageChunks:
                                imageData = imageChunks[0].get("data", "")
                except Exception as extractErr:
                    logger.warning(f"describeImage: on-demand extraction failed: {extractErr}")

            # 3) Direct image file (not a container) - use raw file data
            if not imageData:
                chatService = services.chat
                fileContent = chatService.getFileContent(fileId)
                if fileContent:
                    fileMimeType = fileContent.get("mimeType", "")
                    if fileMimeType.startswith("image/"):
                        imageData = fileContent.get("data", "")
                        mimeType = fileMimeType

            # 4) PDF page rendering: render the requested page as an image via PyMuPDF
            if not imageData:
                chatService = services.chat
                fileInfo = chatService.getFileInfo(fileId) if hasattr(chatService, "getFileInfo") else None
                fileMime = (fileInfo.get("mimeType", "") if fileInfo else "").lower()
                if fileMime == "application/pdf" or (fileInfo and (fileInfo.get("fileName", "") or "").lower().endswith(".pdf")):
                    try:
                        import fitz as _fitz
                        rawContent = chatService.getFileContent(fileId) if not fileContent else fileContent
                        rawData = rawContent.get("data", "") if rawContent else ""
                        if isinstance(rawData, str) and len(rawData) > 100:
                            pdfBytes = _b64.b64decode(rawData)
                        elif isinstance(rawData, bytes):
                            pdfBytes = rawData
                        else:
                            pdfBytes = None
                        if pdfBytes:
                            doc = _fitz.open(stream=pdfBytes, filetype="pdf")
                            targetPage = pageIndex if pageIndex is not None else 0
                            if 0 <= targetPage < len(doc):
                                page = doc[targetPage]
                                pix = page.get_pixmap(dpi=200)
                                imageData = _b64.b64encode(pix.tobytes("png")).decode("ascii")
                                mimeType = "image/png"
                                logger.info("describeImage: rendered PDF page %d as image (%dx%d)", targetPage, pix.width, pix.height)
                            doc.close()
                    except Exception as pdfErr:
                        logger.warning("describeImage: PDF page rendering failed: %s", pdfErr)

            if not imageData:
                chatService = services.chat
                _errFileInfo = chatService.getFileInfo(fileId) if hasattr(chatService, "getFileInfo") else None
                fileName = _errFileInfo.get("fileName", fileId) if _errFileInfo else fileId
                fileMime = _errFileInfo.get("mimeType", "unknown") if _errFileInfo else "unknown"
                return ToolResult(toolCallId="", toolName="describeImage", success=False,
                                  error=f"No image data found in '{fileName}' (type: {fileMime}). "
                                        f"This file likely contains text, not images. Use readFile(fileId=\"{fileId}\") to access its text content.")

            try:
                rawHead = _b64.b64decode(imageData[:32])
                if rawHead[:3] == b"\xff\xd8\xff":
                    mimeType = "image/jpeg"
                elif rawHead[:8] == b"\x89PNG\r\n\x1a\n":
                    mimeType = "image/png"
                elif rawHead[:4] == b"GIF8":
                    mimeType = "image/gif"
                elif rawHead[:4] == b"RIFF" and rawHead[8:12] == b"WEBP":
                    mimeType = "image/webp"
            except Exception as e:
                logger.warning(f"describeImage: MIME detection from base64 header failed for {fileId}: {e}")
            dataUrl = f"data:{mimeType};base64,{imageData}"
            from modules.datamodels.datamodelAi import AiCallRequest, AiCallOptions, OperationTypeEnum as OTE

            _opType = OTE.IMAGE_ANALYSE
            try:
                from modules.datamodels.datamodelFiles import FileItem as _FileItemModel
                from modules.interfaces.interfaceDbManagement import ComponentObjects as _CO
                _fRow = _CO().db._loadRecord(_FileItemModel, fileId)
                if _fRow:
                    _fGet = (lambda k, d=None: _fRow.get(k, d)) if isinstance(_fRow, dict) else (lambda k, d=None: getattr(_fRow, k, d))
                    if bool(_fGet("neutralize", False)):
                        _opType = OTE.NEUTRALIZATION_IMAGE
                        logger.info(f"describeImage: file {fileId} has neutralize=True, using NEUTRALIZATION_IMAGE (internal models only)")
            except Exception as e:
                logger.warning(f"describeImage: neutralize flag check failed for {fileId}: {e}")

            visionRequest = AiCallRequest(
                prompt=prompt,
                options=AiCallOptions(operationType=_opType),
                messages=[{"role": "user", "content": [
                    {"type": "text", "text": prompt},
                    {"type": "image_url", "image_url": {"url": dataUrl}},
                ]}],
            )
            visionResponse = await services.ai.callAi(visionRequest)

            if visionResponse.errorCount > 0:
                return ToolResult(toolCallId="", toolName="describeImage", success=False, error=visionResponse.content)
            return ToolResult(toolCallId="", toolName="describeImage", success=True, data=visionResponse.content)

        except Exception as e:
            return ToolResult(toolCallId="", toolName="describeImage", success=False, error=str(e))

    registry.register(
        "describeImage", _describeImage,
        description="Analyze an image using AI vision. Works with image files and images extracted from PDFs/DOCX/PPTX. Use for OCR, data extraction, and visual analysis.",
        parameters={
            "type": "object",
            "properties": {
                "fileId": {"type": "string", "description": "The file ID containing the image or document with images"},
                "prompt": {"type": "string", "description": "What to look for in the image (default: describe everything)"},
                "pageIndex": {"type": "integer", "description": "Filter images by page index (0-based, for multi-page documents)"},
            },
            "required": ["fileId"],
        },
        readOnly=True,
    )
