"""Gemini batch processing using google-genai SDK."""

import atexit
import json
import math
import mimetypes
import os
import tempfile
import threading
import time
from datetime import datetime
from pathlib import Path
from typing import List, Dict, Any, Optional, Set

from google import genai
from google.genai import types

from .base import BaseBatchManager, BatchJob, BatchStatus


# Module-level temp file tracking for cleanup on exit
_temp_files: Set[str] = set()
_temp_files_lock: Optional[threading.Lock] = threading.Lock()
_initialized = True


def _cleanup_temp_files():
    """Clean up any remaining temp files on module exit."""
    global _temp_files, _temp_files_lock, _initialized
    # Guard against module cleanup race conditions
    if _temp_files_lock is None or _temp_files is None:
        return
    try:
        with _temp_files_lock:
            if _temp_files is None:
                return
            for file_path in list(_temp_files):
                try:
                    os.unlink(file_path)
                except Exception:
                    pass
            _temp_files.clear()
    except Exception:
        pass


def _register_temp_file(file_path: str) -> None:
    """Register a temp file for cleanup."""
    global _temp_files, _temp_files_lock
    if _temp_files_lock is None or _temp_files is None:
        return
    with _temp_files_lock:
        if _temp_files is not None:
            _temp_files.add(file_path)


def _unregister_temp_file(file_path: str) -> None:
    """Unregister a temp file after successful cleanup."""
    global _temp_files, _temp_files_lock
    if _temp_files_lock is None or _temp_files is None:
        return
    with _temp_files_lock:
        if _temp_files is not None:
            _temp_files.discard(file_path)


# Register cleanup only once at module import
atexit.register(_cleanup_temp_files)


class GeminiBatchManager(BaseBatchManager):
    """
    Gemini batch processing manager.

    Uses Google's Batch API for significant cost savings.
    """

    def __init__(
        self,
        api_key: str,
        base_url: Optional[str] = None,
        timeout: Optional[float] = None
    ):
        super().__init__(api_key, base_url, timeout)

        client_kwargs = {"api_key": self.api_key}
        if self.base_url or self.timeout:
            http_options = {}
            if self.base_url:
                http_options["base_url"] = self.base_url
            if self.timeout:
                http_options["timeout"] = math.ceil(self.timeout)
            client_kwargs["http_options"] = types.HttpOptions(**http_options)

        self.client = genai.Client(**client_kwargs)

    def _upload_file(self, file_path: Path) -> str:
        """Upload file and return URI."""
        path = Path(file_path)
        if not path.exists():
            raise FileNotFoundError(f"File not found: {path}")
        if not path.is_file():
            raise ValueError(f"Path is not a file: {path}")

        mime_type = mimetypes.guess_type(str(path))[0] or "application/octet-stream"

        # 透明桥接：直接调用 SDK，让原始异常透传
        with open(path, "rb") as f:
            uploaded = self.client.files.upload(
                file=f,
                config=types.UploadFileConfig(
                    display_name=path.name,
                    mime_type=mime_type
                )
            )

        # Wait for processing (with timeout)
        # 使用与普通层一致的防御性访问，避免 SDK 返回异常对象时崩溃
        max_wait = 300  # 5 minutes
        start_time = time.time()
        last_state = getattr(getattr(uploaded, "state", None), "name", "UNKNOWN")
        while getattr(getattr(uploaded, "state", None), "name", None) == "PROCESSING":
            elapsed = time.time() - start_time
            if elapsed > max_wait:
                raise RuntimeError(
                    f"Gemini file processing timed out after {max_wait}s "
                    f"for '{path}' (file={getattr(uploaded, 'name', None)}, last_state={last_state})."
                )
            time.sleep(1)
            # 透明桥接：状态轮询失败让 SDK 原始异常透传
            uploaded = self.client.files.get(name=uploaded.name)
            last_state = getattr(getattr(uploaded, "state", None), "name", "UNKNOWN")

        final_state = getattr(getattr(uploaded, "state", None), "name", "UNKNOWN")
        if final_state != "ACTIVE":
            raise RuntimeError(
                f"Gemini file upload failed for '{path}' "
                f"(file={getattr(uploaded, 'name', None)}, state={final_state})"
            )

        return uploaded.uri

    def _extract_result_text(self, obj: Dict[str, Any]) -> str:
        """Extract text parts from a Gemini batch result object."""
        response = obj.get("response") or {}
        candidates = response.get("candidates") or []
        if not candidates:
            return ""

        parts = ((candidates[0].get("content") or {}).get("parts")) or []
        texts = []
        for part in parts:
            if isinstance(part, dict):
                text = part.get("text")
            else:
                text = getattr(part, "text", None)
            if text:
                texts.append(str(text))
        return "\n".join(texts)

    def submit(
        self,
        requests: List[Dict[str, Any]],
        model: str,
        **kwargs
    ) -> BatchJob:
        """
        Submit batch job to Gemini.

        Args:
            requests: List of dicts with 'prompt' and optional 'files' keys.
            model: Model name (e.g., 'gemini-2.5-flash').

        Returns:
            BatchJob with Gemini job name.
        """
        # Build JSONL
        batch_requests = []
        for i, req in enumerate(requests):
            prompt = req.get("prompt", "")
            files = req.get("files", [])

            parts = []

            # Upload files
            for file_path in files:
                path = Path(file_path)
                uri = self._upload_file(path)
                mime_type = mimetypes.guess_type(str(path))[0] or "application/octet-stream"
                parts.append({
                    "file_data": {
                        "file_uri": uri,
                        "mime_type": mime_type
                    }
                })

            # Add text prompt
            parts.append({"text": prompt})

            batch_requests.append({
                "key": f"req_{i}",
                "request": {
                    "contents": [{"role": "user", "parts": parts}]
                }
            })

        # Write to temp file and register for cleanup
        with tempfile.NamedTemporaryFile(
            mode="w",
            suffix=".jsonl",
            delete=False,
            encoding="utf-8"
        ) as f:
            for req in batch_requests:
                f.write(json.dumps(req, ensure_ascii=False) + "\n")
            jsonl_path = f.name

        # Register for atexit cleanup in case of process termination
        _register_temp_file(jsonl_path)

        try:
            # Upload batch file (pass file object, not string path, for SDK compatibility)
            with open(jsonl_path, "rb") as f:
                batch_file = self.client.files.upload(
                    file=f,
                    config=types.UploadFileConfig(
                        display_name="batch_input.jsonl",
                        mime_type="application/jsonl"
                    )
                )

            # Create batch job
            job = self.client.batches.create(
                model=model,
                src=batch_file.name,
                config=types.CreateBatchJobConfig(
                    display_name=f"batch_{datetime.now().strftime('%Y%m%d%H%M%S')}"
                )
            )

            return BatchJob(
                id=job.name,
                vendor="gemini",
                model=model,
                status=self._map_status(job.state),
                total_requests=len(requests)
            )

        finally:
            # Clean up temp file
            try:
                os.unlink(jsonl_path)
                _unregister_temp_file(jsonl_path)
            except Exception:
                pass

    def get_status(self, job_id: str) -> BatchJob:
        """Get status of Gemini batch job."""
        job = self.client.batches.get(name=job_id)

        request_counts = getattr(job, "request_counts", None)
        total = (getattr(request_counts, "total", None) or 0) if request_counts else 0
        completed = (getattr(request_counts, "completed", None) or 0) if request_counts else 0

        return BatchJob(
            id=job.name,
            vendor="gemini",
            model=getattr(job, "model", ""),
            status=self._map_status(job.state),
            total_requests=total,
            completed_requests=completed,
            output_file_id=getattr(job, "output_file", None),
        )

    def get_results(self, job_id: str) -> Dict[str, str]:
        """Get results of completed Gemini batch job."""
        job = self.client.batches.get(name=job_id)

        if not hasattr(job, 'output_file') or not job.output_file:
            raise RuntimeError("Job has no output file")

        content = self.client.files.content(name=job.output_file).text

        results = {}
        for line in content.splitlines():
            if not line.strip():
                continue
            try:
                obj = json.loads(line)
                key = obj.get("key", "unknown")
                response_content = self._extract_result_text(obj)
                if response_content:
                    results[key] = response_content
            except (KeyError, json.JSONDecodeError, IndexError, TypeError) as e:
                import warnings
                warnings.warn(f"Failed to parse batch result: {e}")

        return results

    def cancel(self, job_id: str) -> bool:
        """Cancel Gemini batch job."""
        try:
            self.client.batches.cancel(name=job_id)
            return True
        except Exception:
            return False

    def _map_status(self, state) -> BatchStatus:
        """Map Gemini state to BatchStatus."""
        state_name = state if isinstance(state, str) else state.name
        mapping = {
            "JOB_STATE_PENDING": BatchStatus.PENDING,
            "JOB_STATE_RUNNING": BatchStatus.RUNNING,
            "JOB_STATE_SUCCEEDED": BatchStatus.COMPLETED,
            "JOB_STATE_FAILED": BatchStatus.FAILED,
            "JOB_STATE_CANCELLED": BatchStatus.CANCELLED,
            "JOB_STATE_EXPIRED": BatchStatus.EXPIRED,
            "PENDING": BatchStatus.PENDING,
            "RUNNING": BatchStatus.RUNNING,
            "SUCCEEDED": BatchStatus.COMPLETED,
            "COMPLETED": BatchStatus.COMPLETED,
            "FAILED": BatchStatus.FAILED,
            "CANCELLED": BatchStatus.CANCELLED,
            "EXPIRED": BatchStatus.EXPIRED,
        }
        return mapping.get(state_name, BatchStatus.PENDING)
