from __future__ import annotations

import asyncio
import base64
import concurrent.futures
import io
import logging
import os
import time
from typing import Any

from app.core.ai.domain.models import (
    AiModelSpec,
    AiProviderType,
    AiResponse,
    AiStructuredRequest,
    AiTextRequest,
    AiVisionRequest,
)

logger = logging.getLogger("dscons.ai.adapter.antigravity")


class AntigravitySdkAdapter:
    """Primary Enterprise Native Adapter using Google Antigravity SDK."""

    def __init__(
        self,
        api_key: str | None = None,
        timeout_seconds: float = 15.0,
    ) -> None:
        self._api_key = api_key or os.environ.get("GEMINI_API_KEY", "")
        self._timeout = timeout_seconds

    async def generate_text(
        self, request: AiTextRequest, spec: AiModelSpec
    ) -> AiResponse:
        """Executes text generation via Google Antigravity SDK."""
        t0 = time.perf_counter()

        # Assemble prompt
        prompt = ""
        if request.system_prompt:
            prompt += f"SYSTEM INSTRUCTION:\n{request.system_prompt}\n\n"

        if request.messages:
            for m in request.messages:
                role = m.get("role", "user").upper()
                c = m.get("content", "")
                if isinstance(c, list):
                    text_parts = [
                        p.get("text", "")
                        for p in c
                        if isinstance(p, dict) and p.get("type") == "text"
                    ]
                    c = "\n".join(text_parts)
                prompt += f"{role}:\n{c}\n\n"
        elif request.prompt:
            prompt += f"USER:\n{request.prompt}\n"

        content = ""
        try:
            from google.antigravity import Agent, LocalAgentConfig

            config = LocalAgentConfig(
                model=spec.provider_model_id,
                api_key=self._api_key,
            )
            async with Agent(config) as agent:
                resp = await asyncio.wait_for(agent.chat(prompt), timeout=self._timeout)
                content = await resp.text()
        except ImportError:
            # Fallback to official Google GenAI SDK (google-genai)
            from google import genai
            from google.genai import types

            client = genai.Client(api_key=self._api_key or os.environ.get("GEMINI_API_KEY", ""))
            gen_config = types.GenerateContentConfig(
                temperature=request.temperature or 0.2,
                system_instruction=request.system_prompt or None,
            )
            resp = await asyncio.wait_for(
                client.aio.models.generate_content(
                    model=spec.provider_model_id,
                    contents=prompt,
                    config=gen_config,
                ),
                timeout=self._timeout,
            )
            content = resp.text or ""

        latency = (time.perf_counter() - t0) * 1000.0
        return AiResponse(
            content=content,
            model_used=spec.canonical_id,
            provider_used=AiProviderType.ANTIGRAVITY,
            latency_ms=latency,
        )

    async def analyze_vision(
        self, request: AiVisionRequest, spec: AiModelSpec
    ) -> AiResponse:
        """Executes multimodal vision analysis via Google Antigravity SDK or Google GenAI."""
        t0 = time.perf_counter()
        img_bytes = request.image_bytes
        mime_type = request.image_mime or "image/jpeg"

        if not img_bytes and request.data_uri and "," in request.data_uri:
            b64_data = request.data_uri.split(",", 1)[1]
            img_bytes = base64.b64decode(b64_data)

        if not img_bytes and request.file_path and os.path.exists(request.file_path):
            with open(request.file_path, "rb") as f:
                img_bytes = f.read()

        if not img_bytes:
            raise ValueError("No image payload provided for vision analysis.")

        content = ""
        try:
            from google.antigravity import Agent, LocalAgentConfig
            from google.antigravity.types import Image

            img_obj = Image(data=img_bytes, mime_type=mime_type)
            config = LocalAgentConfig(
                model=spec.provider_model_id,
                api_key=self._api_key,
            )
            async with Agent(config) as agent:
                resp = await asyncio.wait_for(
                    agent.chat([request.prompt, img_obj]), timeout=self._timeout
                )
                content = await resp.text()
        except ImportError:
            # Fallback to official Google GenAI SDK (google-genai)
            from google import genai
            from google.genai import types

            client = genai.Client(api_key=self._api_key or os.environ.get("GEMINI_API_KEY", ""))
            part = types.Part.from_bytes(data=img_bytes, mime_type=mime_type)
            gen_config = types.GenerateContentConfig(
                temperature=0.2,
            )
            resp = await asyncio.wait_for(
                client.aio.models.generate_content(
                    model=spec.provider_model_id,
                    contents=[request.prompt or "Phân tích tài liệu hình ảnh", part],
                    config=gen_config,
                ),
                timeout=self._timeout,
            )
            content = resp.text or ""

        latency = (time.perf_counter() - t0) * 1000.0
        return AiResponse(
            content=content,
            model_used=spec.canonical_id,
            provider_used=AiProviderType.ANTIGRAVITY,
            latency_ms=latency,
        )