from __future__ import annotations

import asyncio
import base64
import mimetypes
from pathlib import Path
from typing import Any, Dict, Optional, Tuple, Union

from ..processing.document_to_pdf import convert_to_pdf, ext_from_mime, needs_conversion
from ..utils.exceptions import InferenceError, ValidationError
from ..utils.logger import PerformanceTimer


class MultimodalMixin:
    async def query_multimodal(
        self,
        input_text: str,
        *,
        image_data: Optional[Union[bytes, str, Path, Any]] = None,
        document_path: Optional[Union[str, Path]] = None,
        document_data: Optional[Union[bytes, str]] = None,
        mime_type: Optional[str] = None,
        context: Optional[Dict[str, Any]] = None,
        response_depth: Optional[str] = None,
        analysis_type: str = "comprehensive",
        processing_mode: str = "direct",
        **kwargs,
    ) -> str:
        """
        Multimodal query: accept text + an image or a document.

        Exactly one of: `image_data`, `document_path`, `document_data`.

        By default (`processing_mode="direct"`), attachments are sent directly to Gemini as bytes
        (no local PDF rendering/OCR dependencies). Use `processing_mode="local"` to route documents
        through `GeminiDocumentProcessor` (optional heavy dependencies; better for scanned PDFs).
        """
        # Import from toxo_layer module to avoid circular imports.
        from .toxo_layer import DEPTH_MAX_TOKENS, RESPONSE_DEPTH_INSTRUCTIONS, _should_escalate_depth

        provided = sum(1 for x in (image_data, document_path, document_data) if x is not None)
        if provided != 1:
            raise ValidationError("Provide exactly one of: image_data, document_path, document_data")

        if (getattr(self, "llm_provider", None) or "gemini").lower() != "gemini":
            raise ValidationError("Multimodal querying currently requires provider='gemini' (vision/document support).")
        if not getattr(self, "gemini_client", None):
            raise ValidationError("Gemini client not configured. Call setup_api_key(..., provider='gemini') first.")

        depth = response_depth or getattr(getattr(self.config, "prompt", None), "response_depth", None) or "balanced"
        if depth not in RESPONSE_DEPTH_INSTRUCTIONS:
            depth = "balanced"
        effective_depth = _should_escalate_depth(input_text, depth, getattr(self, "domain", "general"))

        if "max_tokens" not in kwargs and "max_output_tokens" not in kwargs:
            kwargs["max_tokens"] = min(
                getattr(getattr(self.config, "gemini", None), "max_tokens", 8192),
                DEPTH_MAX_TOKENS.get(effective_depth, 768),
            )

        with PerformanceTimer(self.logger, "query_multimodal"):
            try:
                enriched_prompt = await self._process_input(
                    input_text,
                    context=context,
                    response_depth=effective_depth,
                )

                if processing_mode not in ("direct", "local"):
                    raise ValidationError("processing_mode must be 'direct' or 'local'")

                # Default: direct-to-Gemini bytes. Convert DOCX/HTML/etc to PDF when needed.
                if processing_mode == "direct":
                    file_bytes, inferred_mime, ext = self._load_attachment_bytes_and_mime(
                        image_data=image_data,
                        document_path=document_path,
                        document_data=document_data,
                        mime_type=mime_type,
                    )
                    if needs_conversion(inferred_mime, ext):
                        file_bytes = convert_to_pdf(file_bytes, ext=ext, mime_type=inferred_mime)
                        inferred_mime = "application/pdf"
                    response_text = await self.gemini_client.generate_with_file(
                        prompt=enriched_prompt,
                        file_bytes=file_bytes,
                        mime_type=inferred_mime,
                        analysis_type=analysis_type,
                        **kwargs,
                    )
                else:
                    # Optional: route documents through local processing pipeline.
                    if image_data is not None:
                        analysis = await self.gemini_client.analyze_image(
                            image_data=image_data,
                            prompt=enriched_prompt,
                            analysis_type=analysis_type,
                            **kwargs,
                        )
                        response_text = (
                            analysis.get("analysis")
                            if isinstance(analysis, dict) and "analysis" in analysis
                            else (analysis.get("raw_response") if isinstance(analysis, dict) else str(analysis))
                        )
                    else:
                        try:
                            from ..processing.gemini_document_processor import GeminiDocumentProcessor
                        except Exception as e:
                            raise ValidationError(f"GeminiDocumentProcessor not available: {e}")

                        processor = GeminiDocumentProcessor(gemini_client=self.gemini_client, config=self.config)
                        result = await processor.process_document(document_path if document_path is not None else "")

                        insights = getattr(result, "gemini_insights", None)
                        if insights is not None and hasattr(insights, "actionable_insights"):
                            response_text = "\n".join(
                                [
                                    f"Document type: {getattr(insights, 'document_type', 'unknown')}",
                                    *[f"- {x}" for x in (getattr(insights, 'actionable_insights', []) or [])[:8]],
                                ]
                            ).strip()
                        else:
                            response_text = str(result)

                final_response = await self._post_process_response(
                    response_text,
                    context=context,
                    response_depth=effective_depth,
                )
                await self._update_memory(input_text, final_response, context)
                return final_response
            except Exception as e:
                self.logger.error(f"Multimodal query failed: {str(e)}")
                raise InferenceError(f"Failed to process multimodal query: {str(e)}", query=input_text)

    def query_multimodal_sync(self, *args, **kwargs) -> str:
        """Synchronous wrapper for query_multimodal()."""
        return asyncio.run(self.query_multimodal(*args, **kwargs))

    def _decode_data_url_or_base64(self, payload: Union[bytes, str], *, default_mime: str) -> Tuple[bytes, str]:
        """
        Decode:
        - raw bytes => (bytes, default_mime)
        - base64 string => (decoded_bytes, default_mime)
        - data URL => (decoded_bytes, mime_from_header)
        """
        if isinstance(payload, (bytes, bytearray)):
            return bytes(payload), default_mime
        if not isinstance(payload, str):
            raise ValidationError("payload must be bytes or a base64/data URL string")

        s = payload.strip()
        if s.startswith("data:"):
            try:
                header, b64 = s.split(",", 1)
            except ValueError:
                raise ValidationError("Invalid data URL format")
            if ";base64" not in header.lower():
                raise ValidationError("data URL must be base64 encoded")
            mime = header[5:].split(";", 1)[0].strip().lower() or default_mime
            s = b64.strip()
        else:
            mime = default_mime

        try:
            return base64.b64decode(s, validate=True), mime
        except Exception as e:
            raise ValidationError(f"Failed to decode base64 payload: {e}")

    def _load_attachment_bytes_and_mime(
        self,
        *,
        image_data: Optional[Union[bytes, str, Path, Any]],
        document_path: Optional[Union[str, Path]],
        document_data: Optional[Union[bytes, str]],
        mime_type: Optional[str],
    ) -> Tuple[bytes, str, str]:
        """
        Convert the provided multimodal input into (bytes, mime_type, ext) for direct Gemini upload.
        """
        def _norm_ext(e: str) -> str:
            if not e:
                return ""
            e = e.lower().strip()
            return e if e.startswith(".") else f".{e}"

        if image_data is not None:
            # Path-like
            if isinstance(image_data, Path):
                p = image_data
                data = p.read_bytes()
                inferred = mime_type or (mimetypes.guess_type(str(p))[0] or "image/png")
                ext = _norm_ext(p.suffix) or ext_from_mime(inferred)
                return data, inferred, ext
            if isinstance(image_data, str):
                s = image_data.strip()
                # Handle data URLs/base64 first to avoid treating large payloads as filesystem paths.
                if s.startswith("data:"):
                    data, inferred = self._decode_data_url_or_base64(s, default_mime=mime_type or "image/png")
                    ext = ext_from_mime(inferred)
                    return data, inferred, ext
                is_path = False
                # Keep path support for local mode while guarding against OSError on huge strings.
                if len(s) <= 1024:
                    try:
                        is_path = Path(s).exists()
                    except OSError:
                        is_path = False
                if is_path:
                    p = Path(s)
                    data = p.read_bytes()
                    inferred = mime_type or (mimetypes.guess_type(str(p))[0] or "image/png")
                    ext = _norm_ext(p.suffix) or ext_from_mime(inferred)
                    return data, inferred, ext
                data, inferred = self._decode_data_url_or_base64(s, default_mime=mime_type or "image/png")
                ext = ext_from_mime(inferred)
                return data, inferred, ext
            # Data URL or base64
            if isinstance(image_data, (bytes, bytearray)):
                inf = mime_type or "image/png"
                return bytes(image_data), inf, ext_from_mime(inf)
            raise ValidationError("Unsupported image_data type for direct mode. Pass bytes, base64/data URL, or path.")

        # Document path
        if document_path is not None:
            p = Path(document_path)
            if not p.exists():
                raise ValidationError(f"document_path not found: {p}")
            data = p.read_bytes()
            inferred = mime_type or (mimetypes.guess_type(str(p))[0] or "application/octet-stream")
            ext = _norm_ext(p.suffix) or ext_from_mime(inferred)
            return data, inferred, ext

        # Document bytes/base64/data URL
        if document_data is not None:
            default = mime_type or "application/pdf"
            data, inferred = self._decode_data_url_or_base64(document_data, default_mime=default)
            ext = ext_from_mime(inferred)
            return data, inferred, ext

        raise ValidationError("No attachment provided")

