"""
AIGCClient - handles AIGC service requests for Prompt and Agent applications.
"""

import time
import random
import string
import logging
import requests
from typing import Dict, Optional

from .exceptions import AIGCRequestError


logger = logging.getLogger(__name__)


class AIGCClient:
    """
    AIGC service client for handling Prompt and Agent application requests.
    """
    
    def __init__(self, url: str, app_id: str, token: str, aigc_app_id: str, 
                 timeout: int = 300, max_retries: int = 3, sleep: float = 2):
        """
        Initialize AIGC client.
        
        Args:
            url: AIGC service URL
            app_id: Application ID
            token: Authentication token
            aigc_app_id: AIGC application ID
            timeout: Request timeout in seconds (default: 300)
            max_retries: Maximum number of retries for failed requests (default: 3)
            sleep: Sleep time before each request in seconds to avoid rate limiting (default: 2)
                   - Applied before every request (including first request)
                   - For retries, uses exponential backoff: sleep * 2^(attempt-1)
                   - Maximum backoff time is capped at 60 seconds
        """
        self.url = url
        self.app_id = app_id
        self.token = token
        self.aigc_app_id = aigc_app_id
        self.timeout = timeout
        self.max_retries = max_retries
        self.sleep = sleep
    
    def call_prompt_app(self, params: Dict) -> Dict:
        """
        Call Prompt application.
        
        Args:
            params: Parameter dictionary
                - Single parameter: {'query': value}
                - Multiple parameters: {'dynamicColumMap': {key: value, ...}}
        
        Returns:
            dict: Response data with added metadata:
                - judge_time: Request timestamp
                - log_id: Request log ID from response
        
        Raises:
            AIGCRequestError: When request fails
        """
        headers = self._build_prompt_headers()
        body = {
            'aigcAppId': self.aigc_app_id,
            **params
        }
        
        logger.debug(f"Calling Prompt app with params: {params}")
        
        # Record request time
        request_time = time.strftime('%Y-%m-%d %H:%M:%S', time.localtime())
        
        response = self._make_request_with_retry(headers, body)
        
        # Validate response code
        if response.get('code') != '000':
            error_msg = response.get('message', 'Unknown error')
            logger.error(f"Prompt app request failed: {error_msg}")
            raise AIGCRequestError(f"Prompt application returned error: {error_msg}")
        
        # Add metadata to response
        response['_judge_time'] = request_time
        response['_log_id'] = response.get('data', {}).get('logId', '')
        
        logger.debug("Prompt app request successful")
        return response
    
    def call_agent_app(self, params: Dict) -> Dict:
        """
        Call Agent application.
        
        Args:
            params: Parameter dictionary
                - Single parameter: {'content': value} - will be wrapped in message format
                - Multiple parameters: {'content': main_value, 'other_param': value, ...}
                  where 'content' goes to message and others go to args
        
        Returns:
            dict: Response data with added metadata:
                - judge_time: Request timestamp
                - log_id: Conversation ID from response
        
        Raises:
            AIGCRequestError: When request fails
        """
        headers = self._build_agent_headers()
        
        # Build the correct request body format
        body = {
            'aigcAppId': self.aigc_app_id,
            'message': [],
            'args': {}
        }
        
        # Handle parameter formatting
        if 'content' in params:
            # Put content in message array
            body['message'] = [{'content': params['content'], 'content_type': 'text'}]
            
            # Put other parameters in args (if any)
            other_params = {k: v for k, v in params.items() if k != 'content'}
            if other_params:
                body['args'] = other_params
        else:
            # If no 'content' key, treat first parameter as content and rest as args
            param_items = list(params.items())
            if param_items:
                first_key, first_value = param_items[0]
                body['message'] = [{'content': first_value, 'content_type': 'text'}]
                
                # Rest go to args
                if len(param_items) > 1:
                    body['args'] = {k: v for k, v in param_items[1:]}
        
        logger.debug(f"Calling Agent app with formatted body: {body}")
        
        # Record request time
        request_time = time.strftime('%Y-%m-%d %H:%M:%S', time.localtime())
        
        response = self._make_request_with_retry(headers, body)
        
        # Validate response code
        if response.get('code') != 'success':
            error_msg = response.get('message', 'Unknown error')
            logger.error(f"Agent app request failed: {error_msg}")
            raise AIGCRequestError(f"Agent application returned error: {error_msg}")
        
        # Add metadata to response
        response['_judge_time'] = request_time
        # For agent, data is a list, conversation_id is at the first element
        data = response.get('data', [])
        if data and len(data) > 0 and isinstance(data[0], dict):
            response['_log_id'] = data[0].get('conversation_id', '')
        else:
            response['_log_id'] = ''
        
        logger.debug("Agent app request successful")
        return response
    
    def _build_prompt_headers(self) -> Dict[str, str]:
        """
        Build headers for Prompt application request.
        
        Returns:
            dict: Request headers with appId, token, timestamp, and x-g-rid
        """
        return {
            'Content-Type': 'application/json',
            'appId': self.app_id,
            'token': self.token,
            'timestamp': str(int(time.time() * 1000)),  # Millisecond timestamp
            'x-g-rid': self._generate_x_g_rid()
        }
    
    def _build_agent_headers(self) -> Dict[str, str]:
        """
        Build headers for Agent application request.
        
        Returns:
            dict: Request headers with appId, token, and X-TraceId
        """
        return {
            'Content-Type': 'application/json',
            'appId': self.app_id,
            'token': self.token,
            'X-TraceId': self._generate_trace_id()
        }
    
    def _generate_x_g_rid(self) -> str:
        """
        Generate x-g-rid: second-level timestamp + 26-digit random string.
        
        Returns:
            str: Generated x-g-rid (e.g., "1701234567abcdefghijklmnopqrstuvwxyz")
        """
        timestamp = str(int(time.time()))  # Second-level timestamp
        random_str = ''.join(random.choices(string.ascii_lowercase + string.digits, k=26))
        return timestamp + random_str
    
    def _generate_trace_id(self) -> str:
        """
        Generate X-TraceId: 32-digit random string.
        
        Returns:
            str: Generated trace ID (32 characters of lowercase letters and digits)
        """
        return ''.join(random.choices(string.ascii_lowercase + string.digits, k=32))
    
    def _make_request_with_retry(self, headers: Dict[str, str], body: Dict) -> Dict:
        """
        Make HTTP request with retry logic.
        
        Args:
            headers: Request headers
            body: Request body
        
        Returns:
            dict: Response JSON
        
        Raises:
            AIGCRequestError: When all retries fail
        """
        last_exception = None
        
        for attempt in range(self.max_retries):
            try:
                # Sleep before each request to avoid rate limiting
                if self.sleep > 0:
                    if attempt == 0:
                        # For first attempt, use configured sleep time
                        # This ensures interval between consecutive successful requests
                        time.sleep(self.sleep)
                        logger.debug(f"Waited {self.sleep} seconds before request")
                    else:
                        # For retries, use exponential backoff based on configured sleep
                        backoff_time = min(self.sleep * (2 ** (attempt - 1)), 60)  # Max 60 seconds
                        logger.info(f"Waiting {backoff_time:.1f} seconds before retry (attempt {attempt + 1})...")
                        time.sleep(backoff_time)
                
                logger.debug(f"Making request (attempt {attempt + 1}/{self.max_retries})")
                
                response = requests.post(
                    self.url,
                    json=body,
                    headers=headers,
                    timeout=self.timeout
                )
                
                response.raise_for_status()
                return response.json()
                
            except requests.exceptions.Timeout as e:
                last_exception = e
                logger.warning(f"Request timeout (attempt {attempt + 1}/{self.max_retries}): {str(e)}")
                # Backoff is handled at the beginning of next iteration
                    
            except requests.exceptions.RequestException as e:
                last_exception = e
                logger.warning(f"Request failed (attempt {attempt + 1}/{self.max_retries}): {str(e)}")
                # Backoff is handled at the beginning of next iteration
        
        # All retries failed
        error_msg = f"Request failed after {self.max_retries} attempts: {str(last_exception)}"
        logger.error(error_msg)
        raise AIGCRequestError(error_msg)
