from __future__ import annotations

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

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 for *all* formats (image/pdf/etc).
                if processing_mode == "direct":
                    file_bytes, inferred_mime = self._load_attachment_bytes_and_mime(
                        image_data=image_data,
                        document_path=document_path,
                        document_data=document_data,
                        mime_type=mime_type,
                    )
                    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]:
        """
        Convert the provided multimodal input into (bytes, mime_type) for direct Gemini upload.
        """
        if image_data is not None:
            # Path-like
            if isinstance(image_data, Path) or (isinstance(image_data, str) and Path(image_data).exists()):
                p = image_data if isinstance(image_data, Path) else Path(image_data)
                data = p.read_bytes()
                inferred = mime_type or (mimetypes.guess_type(str(p))[0] or "image/png")
                return data, inferred
            # Data URL or base64
            if isinstance(image_data, str):
                data, inferred = self._decode_data_url_or_base64(image_data, default_mime=mime_type or "image/png")
                return data, inferred
            if isinstance(image_data, (bytes, bytearray)):
                return bytes(image_data), (mime_type or "image/png")
            # Last resort: let GeminiClient handle it (e.g., PIL image) by using analyze_image.
            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")
            return data, inferred

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

        raise ValidationError("No attachment provided")

