import os
import socket
import requests
import json
from typing import Union

import lazyllm
from lazyllm import LOG
from lazyllm import ModuleBase, ServerModule
from lazyllm.thirdparty import gradio as gr
from lazyllm.flow import Pipeline


class WebUi:
    """A Gradio-based web UI for managing knowledge base files.

This class provides an interactive UI to create/delete groups, upload files, list files, and perform deletion operations via RESTful APIs. It is designed for rapid integration of file and group management.

Args:
    base_url (str): Base URL of the backend API service.
"""
    def __init__(self, base_url) -> None:
        self.base_url = base_url

    def basic_headers(self, content_type=True):
        """
Generate standard HTTP headers.

Args:
    content_type (bool): Whether to include Content-Type in the headers (default: True).

**Returns:**

- dict: Dictionary of HTTP headers.
"""
        return {
            'accept': 'application/json',
            'Content-Type': 'application/json' if content_type else None,
        }

    def muti_headers(
        self,
    ):
        """
Generates multipart form HTTP request headers.
Used for requests requiring multipart/form-data format such as file uploads.

**Returns:**

- Dict: Returns HTTP request header dictionary containing accept header.
"""
        return {'accept': 'application/json'}

    def post_request(self, url, data):
        """
Send a POST request.

Args:
    url (str): Target request URL.
    data (dict): Request data (will be serialized as JSON).

**Returns:**

- dict: JSON response from the server.
"""
        response = requests.post(
            url, headers=self.basic_headers(), data=json.dumps(data)
        )
        return response.json()

    def get_request(self, url):
        """
Send a GET request.

Args:
    url (str): Target request URL.

**Returns:**

- dict: JSON response from the server.
"""
        response = requests.get(url, headers=self.basic_headers(False))
        return response.json()

    def new_group(self, group_name: str):
        """
Create a new file group.

Args:
    group_name (str): Name of the new group.

**Returns:**

- str: Server message about the creation result.
"""
        response = requests.post(
            f'{self.base_url}/new_group?group_name={group_name}',
            headers=self.basic_headers(True),
        )
        return response.json()['msg']

    def delete_group(self, group_name: str):
        """
Delete a specific file group.

Args:
    group_name (str): Name of the group to delete.

**Returns:**

- str: Server message about the deletion.
"""
        response = requests.post(
            f'{self.base_url}/delete_group?group_name={group_name}',
            headers=self.basic_headers(True),
        )
        return response.json()['msg']

    def list_groups(self):
        """
Gets all knowledge base group list.
Sends request to backend API to get all current knowledge base group information.

**Returns:**

- List: Returns group name list.
"""
        response = requests.get(
            f'{self.base_url}/list_kb_groups', headers=self.basic_headers(False)
        )
        return response.json()['data']

    def upload_files(self, group_name: str, override: bool = True):
        """
Upload files to a specified group.

Args:
    group_name (str): Name of the group.
    override (bool): Whether to override existing files (default: True).

**Returns:**

- Any: Data returned by the backend.
"""
        response = requests.post(
            f'{self.base_url}/upload_files?group_name={group_name}&override={override}',
            headers=self.basic_headers(True),
        )
        return response.json()['data']

    def list_files_in_group(self, group_name: str):
        """
List all files within a specific group.

Args:
    group_name (str): Name of the group.

**Returns:**

- List: List of file information.
"""
        response = requests.get(
            f'{self.base_url}/list_files_in_group?group_name={group_name}&alive=True',
            headers=self.basic_headers(False),
        )
        return response.json()['data']

    def delete_file(self, group_name: str, file_ids: list[str]):
        """
Delete specific files from a group.

Args:
    group_name (str): Name of the group.
    file_ids (List[str]): IDs of files to delete.

**Returns:**

- str: Deletion result message.
"""
        response = requests.post(
            f'{self.base_url}/delete_files_from_group',
            headers=self.basic_headers(True),
            json={'group_name': group_name, 'file_ids': file_ids}
        )
        return response.json()['msg']

    def gr_show_list(self, str_list: list, list_name: Union[str, list]):
        """
Display a list of strings as a Gradio DataFrame.

Args:
    str_list (List): List of strings or rows.
    list_name (Union[str, List]): Column name(s) for the table.

**Returns:**

- gr.DataFrame: Gradio DataFrame component.
"""
        if isinstance(list_name, str):
            headers = ['index', list_name]
            value = [[index, str_list[index]] for index in range(len(str_list))]
        else:
            headers = ['index'] + list_name
            value = [[index] + str_list[index:index + len(list_name)] for index in range(len(str_list))]
        return gr.DataFrame(headers=headers, value=value)

    def create_ui(self):
        """
Builds Gradio interface with multiple tabs, providing the following functionalities:
    - Group List: View all group information
    - Upload Files: Select group and upload files
    - Group File List: View files in specified group
    - Delete Files: Delete specified files from group

**Returns:**

- gr.Blocks: A complete Gradio application instance.
"""
        with gr.Blocks(analytics_enabled=False) as demo:
            with gr.Tabs():
                select_group_list = []

                with gr.TabItem('分组列表'):
                    select_group = self.gr_show_list(
                        self.list_groups(), list_name='group_name'
                    )
                    select_group_list.append(select_group)

                with gr.TabItem('上传文件'):

                    def _upload_files(group_name, files):

                        files_to_upload = [
                            ('files', (os.path.basename(file), open(file, 'rb')))
                            for file in files
                        ]

                        url = f'{self.base_url}/add_files_to_group?group_name={group_name}&override=true'
                        response = requests.post(
                            url, files=files_to_upload, headers=self.muti_headers()
                        )
                        response.raise_for_status()
                        response_data = response.json()
                        gr.Info(str(response_data['msg']))

                        for _, (_, file_obj) in files_to_upload:
                            file_obj.close()

                    select_group = gr.Dropdown(self.list_groups(), label='选择分组')
                    select_group.change(lambda x: x, inputs=select_group, outputs=None)

                    up_files = gr.Files(label='上传文件')
                    up_btn = gr.Button('上传')
                    up_btn.click(
                        _upload_files,
                        inputs=[select_group, up_files],
                        outputs=None,
                    )

                    select_group_list.append(select_group)

                with gr.TabItem('分组文件列表'):
                    def _list_group_files(group_name):
                        file_list = self.list_files_in_group(group_name)
                        values = [[i] + file_list[i][:2] for i in range(len(file_list))]
                        return gr.update(
                            value=values
                        )

                    select_group = gr.Dropdown(self.list_groups(), label='选择分组')
                    show_list = self.gr_show_list([], list_name=['file_id', 'file_name'])
                    select_group.change(
                        fn=_list_group_files, inputs=select_group, outputs=show_list
                    )
                    select_group_list.append(select_group)

                with gr.TabItem('删除文件'):

                    def _list_group_files(group_name):
                        file_list = self.list_files_in_group(group_name)
                        file_list = [','.join(file[:2]) for file in file_list]
                        return gr.update(choices=file_list)

                    select_group = gr.Dropdown(self.list_groups(), label='选择分组')
                    select_file = gr.Dropdown([], label='选择文件')
                    select_group.change(
                        fn=_list_group_files, inputs=select_group, outputs=select_file
                    )
                    delete_btn = gr.Button('删除')

                    def _delete_file(group_name, select_file):
                        file_ids = [select_file.split(',')[0]]
                        gr.Info(self.delete_file(group_name, file_ids))
                        return _list_group_files(group_name)

                    delete_btn.click(
                        fn=_delete_file,
                        inputs=[select_group, select_file],
                        outputs=select_file,
                    )
                    select_group_list.append(select_group)

        return demo


class DocWebModule(ModuleBase):
    """Document Web Interface Module, inherits from ModuleBase, provides web-based document management interface.

Args:
    doc_server (ServerModule): Document server module instance providing backend API support
    title (str, optional): Interface title. Defaults to 'Document Management Demo Terminal'
    port (optional): Service port number or port range. Defaults to ``None`` (uses range 20800-20999)
    history (optional): History record list. Defaults to ``None``
    text_mode (optional): Text display mode. Defaults to ``None`` (uses dynamic mode)
    trace_mode (optional): Trace mode. Defaults to ``None`` (uses refresh mode)


Class Attributes:

    Mode: Mode enumeration class containing:
        - Dynamic: Dynamic mode
        - Refresh: Refresh mode
        - Appendix: Appendix mode

Notes:
    - Requires a valid doc_server instance to work with
    - Automatically tries other ports in range when port conflict occurs
    - Releases resources when service is stopped


Examples:
    >>> import lazyllm
    >>> from lazyllm.tools.rag.web import DocWebModule
    >>> from lazyllm import
    >>> doc_server = ServerModule(url="your_url")
    >>> doc_web = DocWebModule(
    >>>   doc_server=doc_server,
    >>>   title="文档管理演示终端",
    >>>   port=range(20800, 20805)  # 自动寻找可用端口)
    >>> deploy_task = doc_web._get_deploy_tasks()
    >>> deploy_task()  
    >>> print(doc_web.url)
    >>> doc_web.stop()
    """
    class Mode:
        Dynamic = 0
        Refresh = 1
        Appendix = 2

    def __init__(self, doc_server: ServerModule, title='文档管理演示终端', port=None,
                 history=None, text_mode=None, trace_mode=None) -> None:
        super().__init__()
        self.title = title
        self.port = port or range(20800, 20999)
        self.history = history or []
        self.trace_mode = trace_mode if trace_mode else DocWebModule.Mode.Refresh
        self.text_mode = text_mode if text_mode else DocWebModule.Mode.Dynamic
        self.doc_server = doc_server
        self._deploy_flag = lazyllm.once_flag()
        self.api_url = ''
        self.url = ''

    def _prepare(self, query, chat_history):
        if chat_history is None:
            chat_history = []
        return '', chat_history + [[query, None]]

    def _clear_history(self):
        return [], '', ''

    def _work(self):
        if isinstance(self.port, (range, tuple, list)):
            port = self._find_can_use_network_port()
        else:
            port = self.port
            assert self._verify_port_access(port), f'port {port} is occupied'

        self.api_url = self.doc_server._url.rsplit('/', 1)[0]
        self.web_ui = WebUi(self.api_url)
        self.demo = self.web_ui.create_ui()
        self.url = f'http://127.0.0.1:{port}'
        self.broadcast_url = f'http://0.0.0.0:{port}'

        self.demo.queue().launch(server_name='0.0.0.0', server_port=port, prevent_thread_lock=True)
        LOG.success('LazyLLM docwebmodule launched successfully: Running on: '
                    f'{self.broadcast_url}, local URL: {self.url}')

    def _get_deploy_tasks(self):
        return Pipeline(self._work)

    def _get_post_process_tasks(self):
        return Pipeline(self._print_url)

    def wait(self):
        """Block current thread waiting for web service to run.

This method blocks the calling thread until the web service is explicitly stopped.

"""
        self.demo.block_thread()

    def stop(self):
        """Stops the web interface service and releases related resources.

"""
        if self.demo:
            self.demo.close()
            del self.demo
            self.demo, self.url = None, ''

    def _find_can_use_network_port(self):
        for port in self.port:
            if self._verify_port_access(port):
                return port
        raise RuntimeError(
            f'The ports in the range {self.port} are all occupied. '
            'Please change the port range or release the relevant ports.'
        )

    def _print_url(self):
        lazyllm.LOG.success(f'LazyLLM DocWebModule launched successfully: Running on local URL: {self.url}')

    def _verify_port_access(self, port):
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
            result = s.connect_ex(('127.0.0.1', port))
            return result != 0

    def __repr__(self):
        return lazyllm.make_repr('Module', 'DocWebModule')
