"""Qwen batch processing using OpenAI-compatible Batch API."""

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

from openai import OpenAI

from .base import BaseBatchManager, BatchJob, BatchStatus


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

# Use lambda with locals() to avoid NameError during module cleanup
atexit.register(lambda: _cleanup_temp_files_safe())


def _cleanup_temp_files_safe():
    """Clean up temp files safely, handling module cleanup edge cases."""
    global _temp_files, _temp_files_lock
    # Guard against None during interpreter shutdown
    lock = _temp_files_lock
    files = _temp_files
    if lock is None or files is None:
        return
    try:
        with lock:
            if _temp_files is not None:
                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)


class QwenBatchManager(BaseBatchManager):
    """
    Qwen batch processing manager.

    Uses Dashscope's OpenAI-compatible Batch API for 50% cost savings.
    """

    DEFAULT_BASE_URL = "https://dashscope.aliyuncs.com/compatible-mode/v1"

    def _validate_file(self, file_path: Path) -> Path:
        """Validate that path exists, is a file, and return resolved path."""
        path = Path(file_path)
        if not path.exists():
            raise FileNotFoundError(f"File not found: {file_path}")
        if not path.is_file():
            raise ValueError(f"Path is not a file: {file_path}")
        return path.resolve()

    def __init__(
        self,
        api_key: str,
        base_url: Optional[str] = None,
        timeout: Optional[float] = None
    ):
        super().__init__(api_key, base_url or self.DEFAULT_BASE_URL, timeout)
        self.client = OpenAI(
            api_key=self.api_key,
            base_url=self.base_url,
            timeout=self.timeout
        )
    
    def submit(
        self,
        requests: List[Dict[str, Any]],
        model: str,
        **kwargs
    ) -> BatchJob:
        """
        Submit batch job to Qwen.
        
        Args:
            requests: List of dicts with 'prompt' and optional 'files' keys.
            model: Model name (e.g., 'qwen-plus').
        
        Returns:
            BatchJob with Qwen job ID.
        """
        # Build JSONL file
        batch_requests = []
        for i, req in enumerate(requests):
            prompt = req.get("prompt", "")
            files = req.get("files", [])
            
            messages = []
            uploaded_file_ids: List[str] = []

            # Handle files (upload via File API)
            try:
                for file_path in files:
                    path = self._validate_file(Path(file_path))
                    with open(path, "rb") as f:
                        uploaded = self.client.files.create(file=f, purpose="file-extract")
                    uploaded_file_ids.append(uploaded.id)
                    messages.append({
                        "role": "system",
                        "content": f"fileid://{uploaded.id}"
                    })
            except Exception:
                # Clean up already-uploaded files on failure to avoid resource leaks
                for fid in uploaded_file_ids:
                    try:
                        self.client.files.delete(fid)
                    except Exception:
                        pass
                raise
            
            messages.append({"role": "user", "content": prompt})
            
            batch_requests.append({
                "custom_id": f"req_{i}",
                "method": "POST",
                "url": "/v1/chat/completions",
                "body": {
                    "model": model,
                    "messages": messages
                }
            })
        
        # Write to temp file
        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 cleanup in case of process crash
        _register_temp_file(jsonl_path)

        try:
            # Upload batch file
            with open(jsonl_path, "rb") as f:
                batch_file = self.client.files.create(file=f, purpose="batch")

            # Create batch job
            job = self.client.batches.create(
                input_file_id=batch_file.id,
                endpoint="/v1/chat/completions",
                completion_window="24h"
            )

            return BatchJob(
                id=job.id,
                vendor="qwen",
                model=model,
                status=self._map_status(job.status),
                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 Qwen batch job."""
        job = self.client.batches.retrieve(job_id)
        
        return BatchJob(
            id=job.id,
            vendor="qwen",
            model="",
            status=self._map_status(job.status),
            completed_requests=(job.request_counts.completed or 0) if job.request_counts else 0,
            total_requests=(job.request_counts.total or 0) if job.request_counts else 0,
            output_file_id=job.output_file_id,
            error=str(job.errors) if job.errors else None
        )
    
    def get_results(self, job_id: str) -> Dict[str, str]:
        """Get results of completed Qwen batch job."""
        job = self.client.batches.retrieve(job_id)
        
        if not job.output_file_id:
            raise RuntimeError("Job has no output file")
        
        content = self.client.files.content(job.output_file_id).text
        
        results = {}
        for line in content.splitlines():
            if not line.strip():
                continue
            try:
                obj = json.loads(line)
                custom_id = obj.get("custom_id", "unknown")
                # Safely extract response content with nested access protection
                response = obj.get("response", {})
                body = response.get("body", {})
                choices = body.get("choices", [])
                if not choices:
                    import warnings
                    warnings.warn(f"No choices in batch result for custom_id {custom_id}")
                    continue
                message = choices[0].get("message", {})
                response_content = message.get("content", "")
                if response_content:
                    results[custom_id] = 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 Qwen batch job."""
        try:
            self.client.batches.cancel(job_id)
            return True
        except Exception:
            return False
    
    def _map_status(self, status: str) -> BatchStatus:
        """Map Qwen status to BatchStatus."""
        mapping = {
            "validating": BatchStatus.VALIDATING,
            "in_progress": BatchStatus.RUNNING,
            "completed": BatchStatus.COMPLETED,
            "failed": BatchStatus.FAILED,
            "cancelled": BatchStatus.CANCELLED,
            "expired": BatchStatus.EXPIRED,
        }
        return mapping.get(status, BatchStatus.PENDING)
