import os, math, time, requests
from tqdm import tqdm
from collections import OrderedDict
from concurrent.futures import ThreadPoolExecutor, as_completed
from .config import (
    VALENCE_API_KEY,
    VALENCE_API_BASE_URL,
    VALENCE_WEBSOCKET_URL,
    VALENCE_DISCRETE_URL,
    VALENCE_ASYNCH_URL
)
from .logger import get_logger
from .exceptions import ValenceSDKException, UploadError, PredictionError, AudioTooShortError, AudioTooLongError, FileSizeLimitExceededError
from .rate_limit import RateLimitAPI
from .streaming import StreamingAPI

logger = get_logger()


def _check_audio_validation_error(response):
    """Check if response contains audio validation errors and raise appropriate exception."""
    if response.status_code == 400:
        try:
            error_data = response.json()
            error_code = error_data.get("error_code")
            if error_code == "AUDIO_TOO_SHORT":
                raise AudioTooShortError(
                    error_data.get("error", "Audio file is too short"),
                    min_duration=error_data.get("min_duration_seconds"),
                    actual_duration=error_data.get("actual_duration_seconds")
                )
            elif error_code in ("AUDIO_TOO_LONG", "FILE_TOO_LARGE"):
                raise AudioTooLongError(
                    error_data.get("error", "Audio file exceeds maximum allowed duration"),
                    max_duration=error_data.get("max_duration_seconds"),
                    actual_duration=error_data.get("actual_duration_seconds")
                )
        except ValueError:
            pass  # Response is not JSON, let raise_for_status handle it
    response.raise_for_status()


class DiscreteClient:
    """Client for discrete (short) audio processing."""

    def __init__(self, headers, base_url):
        self.headers = headers
        self.base_url = base_url

    def emotions(self, file_path: str = None, audio_array: list = None, model: str = "4emotions"):
        """Get emotions for discrete (short) audio files.

        Args:
            file_path (str): Path to the audio file (mutually exclusive with audio_array)
            audio_array (list): Audio data as array (mutually exclusive with file_path)
            model (str): Model type ('4emotions' or '7emotions')

        Returns:
            dict: Emotion prediction results

        Raises:
            ValueError: If both or neither file_path and audio_array are provided
            FileSizeLimitExceededError: If file size exceeds 6MB limit
        """
        if file_path and audio_array:
            raise ValueError("Provide either file_path or audio_array, not both")
        if not file_path and not audio_array:
            raise ValueError("Must provide either file_path or audio_array")

        url = f"{self.base_url}/v1/discrete/emotion"
        params = {"model": model}

        if file_path:
            # Validate file size (6MB limit for discrete API)
            file_size_bytes = os.path.getsize(file_path)
            file_size_mb = file_size_bytes / (1024 * 1024)
            max_size_mb = 6.0

            if file_size_mb > max_size_mb:
                error_msg = f"File size ({file_size_mb:.2f} MB) exceeds discrete API maximum ({max_size_mb} MB)"
                logger.error(error_msg)
                raise FileSizeLimitExceededError(
                    error_msg,
                    max_size_mb=max_size_mb,
                    actual_size_mb=file_size_mb
                )

            # File upload method
            logger.info(f"Getting emotions for discrete audio: {file_path} with model={model}")
            with open(file_path, "rb") as f:
                files = {"file": f}
                response = requests.post(url, headers=self.headers, files=files, params=params)
                _check_audio_validation_error(response)
                return response.json()
        else:
            # JSON payload method
            logger.info(f"Getting emotions for audio array with model={model}")
            payload = {"payload": [audio_array]}
            response = requests.post(url, json=payload, headers=self.headers, params=params)
            _check_audio_validation_error(response)
            return response.json()

class AsyncClient:
    """Client for async (long) audio processing."""

    def __init__(self, headers, base_url, part_size=5*1024*1024, show_progress=True, max_threads=3, comprehensive_output=False):
        self.headers = headers
        self.base_url = base_url
        self.part_size = part_size
        self.show_progress = show_progress
        self.max_threads = max_threads
        self.comprehensive_output = comprehensive_output
        self._rate_limit_api = None

    def _get_rate_limit_api(self):
        """Lazy initialization of RateLimitAPI."""
        if self._rate_limit_api is None:
            self._rate_limit_api = RateLimitAPI(self.headers, self.base_url)
        return self._rate_limit_api

    def _validate_file_size(self, file_path: str):
        """Validate file size against rate limit policy.

        Args:
            file_path (str): Path to the audio file

        Raises:
            FileSizeLimitExceededError: If file size exceeds the rate limit policy maximum
        """
        file_size_bytes = os.path.getsize(file_path)
        file_size_mb = file_size_bytes / (1024 * 1024)

        try:
            # Get rate limit status to check max_audio_size_mb
            rate_limit_status = self._get_rate_limit_api().get_status()
            limits = rate_limit_status.get("limits", {})
            max_size_mb = limits.get("max_audio_size_mb")

            # If no policy is set, default to 1GB (1024 MB)
            if max_size_mb is None:
                max_size_mb = 1024
                logger.debug("No rate limit policy found, using default maximum of 1GB")

            if file_size_mb > max_size_mb:
                error_msg = f"File size ({file_size_mb:.2f} MB) exceeds rate limit maximum ({max_size_mb} MB)"
                logger.error(error_msg)
                raise FileSizeLimitExceededError(
                    error_msg,
                    max_size_mb=max_size_mb,
                    actual_size_mb=file_size_mb
                )

            logger.info(f"File size validation passed: {file_size_mb:.2f} MB / {max_size_mb} MB")

        except FileSizeLimitExceededError:
            # Re-raise file size exceeded errors
            raise
        except Exception as e:
            # If rate limit check fails, log warning but don't block upload
            # This ensures backward compatibility if rate limit API is unavailable
            logger.warning(f"Could not validate file size against rate limits: {e}")
            logger.info("Proceeding with upload without rate limit validation")
    
    def upload(self, file_path: str):
        """Upload async (long) audio files using multipart upload.

        Args:
            file_path (str): Path to the audio file

        Returns:
            str: Request ID for tracking the upload

        Raises:
            FileSizeLimitExceededError: If file size exceeds rate limit maximum
        """
        # Validate file size against rate limit policy before upload
        self._validate_file_size(file_path)

        file_size = os.path.getsize(file_path)
        part_count = math.ceil(file_size / self.part_size)
        init_url = f"{self.base_url}/v1/asynch/emotion/upload/initiate"
        params = {"file_name": os.path.basename(file_path), "part_count": part_count}
        response = requests.post(init_url, headers=self.headers, params=params)
        _check_audio_validation_error(response)
        data = response.json()
        upload_id, request_id, presigned_urls = data["upload_id"], data["request_id"], data["presigned_urls"]

        parts = []
        with open(file_path, "rb") as f, ThreadPoolExecutor(max_workers=self.max_threads) as executor:
            futures = []
            for part in presigned_urls:
                part_number = part["part_number"]
                data = f.read(self.part_size)
                future = executor.submit(self._upload_part, part["url"], data, part_number)
                futures.append(future)
            if self.show_progress:
                for f in tqdm(as_completed(futures), total=len(futures), desc="Uploading parts"):
                    part_result = f.result()
                    if part_result:
                        parts.append(part_result)
            else:
                for f in as_completed(futures):
                    part_result = f.result()
                    if part_result:
                        parts.append(part_result)

        complete_payload = OrderedDict([
            ("request_id", request_id),
            ("upload_id", upload_id),
            ("parts", sorted(parts, key=lambda x: x['PartNumber']))
        ])
        comp_url = f"{self.base_url}/v1/asynch/emotion/upload/complete"
        complete_resp = requests.post(comp_url, json=complete_payload, headers=self.headers)
        _check_audio_validation_error(complete_resp)
        logger.info(f"Async upload completed. Request ID: {request_id}")
        return request_id
    
    def _upload_part(self, url, data, part_number):
        """Upload a single part of a multipart upload."""
        try:
            headers = {"Content-Length": str(len(data))}
            resp = requests.put(url, data=data, headers=headers)
            resp.raise_for_status()
            return {"ETag": resp.headers["ETag"], "PartNumber": part_number}
        except Exception as e:
            logger.error(f"Upload failed on part {part_number}")
            raise UploadError(e)
    
    def emotions(self, request_id: str, max_attempts: int = 20, interval_seconds: int = 5):
        """Get emotion prediction results for async audio.

        Args:
            request_id (str): Request ID from upload method
            max_attempts (int): Maximum polling attempts
            interval_seconds (int): Seconds between polling attempts

        Returns:
            dict: Emotion prediction results with timeline data
        """
        url = f"{self.base_url}/v1/asynch/emotion/status/{request_id}"
        params = {"comprehensive_output": "true"} if self.comprehensive_output else {}
        for _ in range(max_attempts):
            time.sleep(interval_seconds)
            resp = requests.get(url, headers=self.headers, params=params)
            if resp.status_code == 200:
                result = resp.json()
                status = result.get("status")
                if status == "completed":
                    logger.info(f"Emotions retrieved for request ID: {request_id}")
                    return result
                elif status in ["processing", "upload_completed", "initiated"]:
                    logger.info(f"Request {request_id} status: {status}, continuing to poll...")
                    continue
                elif status == "failed":
                    raise PredictionError(f"Request {request_id} failed during processing")
        raise PredictionError("Failed to fetch prediction after retries.")

    def get_timeline(self, request_id: str):
        """Get timeline data from completed async emotion prediction.

        Args:
            request_id (str): Request ID from upload method

        Returns:
            list: List of emotion predictions with timestamps
        """
        result = self.emotions(request_id, max_attempts=1)
        if "emotions" in result:
            return result["emotions"]
        return []

    def get_emotion_at_time(self, request_id: str, time_seconds: float):
        """Get emotion prediction at a specific time.

        Args:
            request_id (str): Request ID from upload method
            time_seconds (float): Time in seconds to get emotion for

        Returns:
            dict: Emotion prediction at specified time, or None if not found
        """
        timeline = self.get_timeline(request_id)
        for emotion_data in timeline:
            if emotion_data["start_time"] <= time_seconds < emotion_data["end_time"]:
                return emotion_data
        return None

    def get_dominant_emotion(self, request_id: str):
        """Get the most frequently occurring emotion in the timeline.

        Args:
            request_id (str): Request ID from upload method

        Returns:
            str: The dominant emotion across the timeline
        """
        counts = self.emotion_counts(request_id)
        return max(counts, key=counts.get) if counts else None

    def majority_emotion(self, request_id: str):
        """Alias for get_dominant_emotion. Get the most frequently occurring emotion.

        Args:
            request_id (str): Request ID from upload method

        Returns:
            str: The dominant emotion across the timeline
        """
        return self.get_dominant_emotion(request_id)

    def emotion_counts(self, request_id: str):
        """Get counts of each emotion in the timeline.

        Args:
            request_id (str): Request ID from upload method

        Returns:
            dict: Dictionary mapping emotion names to their occurrence counts
                  (e.g., {"happy": 10, "sad": 3, "angry": 8, "neutral": 9})
        """
        timeline = self.get_timeline(request_id)
        if not timeline:
            return {}

        counts = {}
        for emotion_data in timeline:
            emotion = emotion_data["emotion"]
            counts[emotion] = counts.get(emotion, 0) + 1

        return counts

class ValenceClient:
    """Main client for Valence API with nested discrete and async clients."""

    def __init__(
        self,
        api_key: str = None,
        base_url: str = None,
        websocket_url: str = None,
        part_size=5*1024*1024,
        show_progress=True,
        max_threads=3,
        comprehensive_output=False
    ):
        """Initialize the Valence API client.

        Args:
            api_key: API key for authentication (or set VALENCE_API_KEY env var)
            base_url: Base URL for API endpoints (default: https://demo.getvalenceai.com)
            websocket_url: WebSocket URL for streaming (default: wss://demo.getvalenceai.com)
            part_size: Size of parts for multipart upload (default: 5MB)
            show_progress: Show progress bar for uploads (default: True)
            max_threads: Max concurrent threads for uploads (default: 3)
            comprehensive_output: Include all_predictions in async emotion responses (default: False)
        """
        # Get API key from parameter, environment, or fail
        self.api_key = api_key or VALENCE_API_KEY
        if not self.api_key:
            raise ValenceSDKException("API key not provided and not set in environment.")

        # Get URLs with priority: parameter > env var > default
        self.base_url = base_url or VALENCE_API_BASE_URL
        self.websocket_url = websocket_url or VALENCE_WEBSOCKET_URL

        self.headers = {"x-api-key": self.api_key}

        # Initialize nested clients
        self.discrete = DiscreteClient(self.headers, self.base_url)
        self.asynch = AsyncClient(self.headers, self.base_url, part_size, show_progress, max_threads, comprehensive_output)
        self.rate_limit = RateLimitAPI(self.headers, self.base_url)
        self.streaming = StreamingAPI(self.api_key, self.websocket_url)

        