import requests
import lazyllm
from typing import Tuple, List, Dict, Union
from urllib.parse import urljoin
from ..base import OnlineChatModuleBase, OnlineEmbeddingModuleBase, OnlineMultiModalBase
from lazyllm.components.formatter import encode_query_with_filepaths
from lazyllm.components.utils.file_operate import bytes_to_file
from ..fileHandler import FileHandlerBase


class SiliconFlowModule(OnlineChatModuleBase, FileHandlerBase):
    """SiliconFlow module, inherits from OnlineChatModuleBase and FileHandlerBase.

Provides large language model chat capabilities via the SiliconFlow platform, supports multiple models (including vision-language models), and includes file handling functionality.

Args:
    base_url (str, optional): Base API URL, defaults to "https://api.siliconflow.cn/v1/"
    model (str, optional): Model name to use, defaults to "Qwen/QwQ-32B"
    api_key (str, optional): API key, defaults to lazyllm.config['siliconflow_api_key']
    stream (bool, optional): Whether to enable streaming output, defaults to True
    return_trace (bool, optional): Whether to return trace information, defaults to False
    **kwargs: Other model parameters
"""
    VLM_MODEL_PREFIX = ['Qwen/Qwen2.5-VL-72B-Instruct', 'Qwen/Qwen3-VL-30B-A3B-Instruct', 'deepseek-ai/deepseek-vl2',
                        'Qwen/Qwen3-VL-30B-A3B-Thinking', 'THUDM/GLM-4.1V-9B-Thinking']

    def __init__(self, base_url: str = 'https://api.siliconflow.cn/v1/', model: str = 'Qwen/QwQ-32B',
                 api_key: str = None, stream: bool = True, return_trace: bool = False, **kwargs):
        OnlineChatModuleBase.__init__(self, model_series='SILICONFLOW',
                                      api_key=api_key or lazyllm.config['siliconflow_api_key'],
                                      base_url=base_url, model_name=model, stream=stream,
                                      return_trace=return_trace, **kwargs)
        FileHandlerBase.__init__(self)
        if stream:
            self._model_optional_params['stream'] = True

    def _get_system_prompt(self):
        return 'You are an intelligent assistant provided by SiliconFlow. You are a helpful assistant.'

    def _validate_api_key(self):
        """Validate API Key by sending a minimal request"""
        try:
            # SiliconFlow validates API key using a minimal chat request
            models_url = urljoin(self._base_url, 'models')
            response = requests.get(models_url, headers=self._header, timeout=10)
            return response.status_code == 200
        except Exception:
            return False


class SiliconFlowEmbedding(OnlineEmbeddingModuleBase):
    """SiliconFlow embedding module, inherits from OnlineEmbeddingModuleBase.

Provides text embedding functionality via the SiliconFlow platform, converting text into vector representations.

Args:
    embed_url (str, optional): Embedding API URL, defaults to "https://api.siliconflow.cn/v1/embeddings"
    embed_model_name (str, optional): Name of the embedding model to use, defaults to "BAAI/bge-large-zh-v1.5"
    api_key (str, optional): API key, defaults to lazyllm.config['siliconflow_api_key']
    batch_size (int, optional): Batch size for processing, defaults to 16
    **kw: Additional embedding module parameters
"""
    def __init__(self, embed_url: str = 'https://api.siliconflow.cn/v1/embeddings',
                 embed_model_name: str = 'BAAI/bge-large-zh-v1.5', api_key: str = None,
                 batch_size: int = 16, **kw):
        super().__init__('SILICONFLOW', embed_url, api_key or lazyllm.config['siliconflow_api_key'],
                         embed_model_name, batch_size=batch_size, **kw)


class SiliconFlowReranking(OnlineEmbeddingModuleBase):
    """SiliconFlow reranking module, inherits from OnlineEmbeddingModuleBase.

Provides text reranking functionality via the SiliconFlow platform, reordering a list of documents based on their relevance to a given query.

Args:
    embed_url (str, optional): Reranking API URL, defaults to "https://api.siliconflow.cn/v1/rerank"
    embed_model_name (str, optional): Name of the reranking model to use, defaults to "BAAI/bge-reranker-v2-m3"
    api_key (str, optional): API key, defaults to lazyllm.config['siliconflow_api_key']
    **kw: Additional reranking module parameters
Returns:
    List[Tuple]: A list of reranking results, each containing 'index' and 'relevance_score'.
"""
    def __init__(self, embed_url: str = 'https://api.siliconflow.cn/v1/rerank',
                 embed_model_name: str = 'BAAI/bge-reranker-v2-m3', api_key: str = None, **kw):
        super().__init__('SILICONFLOW', embed_url, api_key or lazyllm.config['siliconflow_api_key'],
                         embed_model_name, **kw)

    def _encapsulated_data(self, query: str, documents: List[str], top_n: int, **kwargs) -> Dict:
        json_data = {
            'model': self._embed_model_name,
            'query': query,
            'documents': documents,
            'top_n': top_n
        }
        if len(kwargs) > 0:
            json_data.update(kwargs)
        return json_data

    def _parse_response(self, response: Dict, input: Union[List, str]) -> List[Tuple]:
        results = response.get('results', [])
        return [(result['index'], result['relevance_score']) for result in results]


class SiliconFlowTextToImageModule(OnlineMultiModalBase):
    """SiliconFlow Text-to-Image module, inherits from OnlineMultiModalBase.

Provides text-to-image generation functionality based on SiliconFlow, supports generating images from text descriptions.

Args:
    api_key (str, optional): API key, defaults to configured siliconflow_api_key
    model_name (str, optional): Model name, defaults to "Qwen/Qwen-Image"
    base_url (str, optional): Base API URL, defaults to "https://api.siliconflow.cn/v1/"
    return_trace (bool, optional): Whether to return trace information, defaults to False
    **kwargs: Other model parameters
"""
    MODEL_NAME = 'Qwen/Qwen-Image'

    def __init__(self, api_key: str = None, model_name: str = None,
                 base_url: str = 'https://api.siliconflow.cn/v1/',
                 return_trace: bool = False, **kwargs):
        OnlineMultiModalBase.__init__(self, model_series='SiliconFlow',
                                      api_key=api_key or lazyllm.config['siliconflow_api_key'],
                                      model_name=model_name or SiliconFlowTextToImageModule.MODEL_NAME,
                                      return_trace=return_trace, **kwargs)
        self._endpoint = 'images/generations'
        self._base_url = base_url

    def _make_request(self, endpoint, payload, timeout=180):
        url = f'{self._base_url}{endpoint}'
        try:
            response = requests.post(url, headers=self._header, json=payload, timeout=timeout)
            response.raise_for_status()
            return response.json()
        except Exception as e:
            lazyllm.LOG.error(f'API request failed: {str(e)}')
            raise

    def _forward(self, input: str = None, size: str = '1024x1024', **kwargs):
        payload = {
            'model': self._model_name,
            'prompt': input
        }
        payload.update(kwargs)

        result = self._make_request(self._endpoint, payload)

        image_urls = [item['url'] for item in result['data']]

        image_files = []
        for url in image_urls:
            img_response = requests.get(url, timeout=180)
            if img_response.status_code == 200:
                image_files.append(img_response.content)
            else:
                raise Exception(f'Failed to download image from {url}')

        file_paths = bytes_to_file(image_files)

        return encode_query_with_filepaths(None, file_paths)


class SiliconFlowTTSModule(OnlineMultiModalBase):
    """SiliconFlow Text-to-Speech module, inherits from OnlineMultiModalBase.

Provides text-to-speech (TTS) functionality based on SiliconFlow, supports converting text to audio files.

Args:
    api_key (str, optional): API key, defaults to configured siliconflow_api_key
    model_name (str, optional): Model name, defaults to "fnlp/MOSS-TTSD-v0.5"
    base_url (str, optional): Base API URL, defaults to "https://api.siliconflow.cn/v1/"
    return_trace (bool, optional): Whether to return trace information, defaults to False
    **kwargs: Other model parameters
"""
    MODEL_NAME = 'fnlp/MOSS-TTSD-v0.5'

    def __init__(self, api_key: str = None, model_name: str = None,
                 base_url: str = 'https://api.siliconflow.cn/v1/',
                 return_trace: bool = False, **kwargs):
        OnlineMultiModalBase.__init__(self, model_series='SiliconFlow',
                                      api_key=api_key or lazyllm.config['siliconflow_api_key'],
                                      model_name=model_name or SiliconFlowTTSModule.MODEL_NAME,
                                      return_trace=return_trace, **kwargs)
        self._endpoint = 'audio/speech'
        self._base_url = base_url

    def _make_binary_request(self, endpoint, payload, timeout=180):
        url = f'{self._base_url}{endpoint}'
        try:
            response = requests.post(url, headers=self._header, json=payload, timeout=timeout)
            response.raise_for_status()
            return response.content
        except Exception as e:
            lazyllm.LOG.error(f'API request failed: {str(e)}')
            raise

    def _forward(self, input: str = None, response_format: str = 'mp3',
                 sample_rate: int = 44100, speed: float = 1.0,
                 voice: str = None, references=None, out_path: str = None, **kwargs):

        if not voice:
            if self._model_name == 'fnlp/MOSS-TTSD-v0.5':
                voice = 'fnlp/MOSS-TTSD-v0.5:alex'
            elif self._model_name == 'FunAudioLLM/CosyVoice2-0.5B':
                voice = 'FunAudioLLM/CosyVoice2-0.5B:alex'
            else:
                raise ValueError(
                    f'Default voice is only supported for models "fnlp/MOSS-TTSD-v0.5" and '
                    f'"FunAudioLLM/CosyVoice2-0.5B". For model "{self._model_name}", '
                    f'please provide a valid voice parameter.')
        payload = {
            'model': self._model_name,
            'input': input,
            'response_format': response_format,
            'sample_rate': sample_rate,
            'speed': speed,
            'voice': voice
        }

        if references:
            payload['references'] = references

        payload.update(kwargs)
        audio_content = self._make_binary_request(self._endpoint, payload, timeout=180)
        file_path = bytes_to_file([audio_content])[0]

        if out_path:
            with open(file_path, 'rb') as src, open(out_path, 'wb') as dst:
                dst.write(src.read())
            file_path = out_path

        result = encode_query_with_filepaths(None, [file_path])

        return result
