#!/usr/bin/env python3
# coding: utf-8
# Copyright 2024 Huawei Technologies Co., Ltd
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <https://www.gnu.org/licenses/>.
# ===========================================================================

import fnmatch

import os
import platform
import shutil

from ansible.module_utils.path_manager import TmpPath
from ansible.module_utils import common_info, common_utils

CONF_DIR = "conf"
EXAMPLES_DIR = "examples"
KUBERNETES_DEPLOY_SCRIPTS = "kubernetes_deploy_scripts"
CONFIG_JSON = "config.json"
CONFIG_P_JSON = "config_p.json"
CONFIG_D_JSON = "config_d.json"


class ExtractMindieDeployer:

    def __init__(self, module):
        self.module = module
        self.resources_dir = os.path.expanduser(module.params["resources_dir"])
        self.arch = platform.machine()
        self.mindie_unzip_path = os.path.join(self.resources_dir, "run_from_mindie")
        self.mindie_service_path = os.path.join(self.mindie_unzip_path, "mindie_service")
        self.mindie_llm_path = os.path.join(self.mindie_unzip_path, "mindie_llm")
        self.mindie_deploy_path = os.path.join(TmpPath.ROOT, "mindie_pd")
        self.messages = []
        self._is_whl_version = False

    @staticmethod
    def _find_data_dir(base_dir, prefix):
        for item in os.listdir(base_dir):
            if item.startswith(prefix) and item.endswith(".data") and os.path.isdir(os.path.join(base_dir, item)):
                return os.path.join(base_dir, item)
        return None

    def find_package(self, pattern):
        for root, dirs, files in os.walk(self.resources_dir):
            for file in files:
                if fnmatch.fnmatch(file, pattern):
                    return os.path.join(root, file)
        return None

    def unzip_package(self):
        if os.path.exists(self.mindie_unzip_path):
            shutil.rmtree(self.mindie_unzip_path)
        os.makedirs(self.mindie_unzip_path, mode=0o750)

        llm_whl = self.find_package("mindie_llm-*linux_{}*.whl".format(self.arch))
        motor_whl = self.find_package("mindie_motor-*linux_{}*.whl".format(self.arch))

        if llm_whl and motor_whl:
            self._is_whl_version = True
            self._unzip_whl_packages(llm_whl, motor_whl)
        else:
            self._is_whl_version = False
            self._unzip_run_packages()
            self.create_config_json()

    def create_config_json(self):
        src_config = os.path.join(self.mindie_llm_path, CONF_DIR, CONFIG_JSON)
        if os.path.exists(src_config):
            for config_name in [CONFIG_JSON, CONFIG_P_JSON, CONFIG_D_JSON]:
                dest_config = os.path.join(self.mindie_service_path, EXAMPLES_DIR,
                                           KUBERNETES_DEPLOY_SCRIPTS, CONF_DIR, config_name)
                shutil.copy2(src_config, dest_config)
                os.chmod(dest_config, 0o640)

    def extract_conf(self):
        if os.path.exists(self.mindie_deploy_path):
            shutil.rmtree(self.mindie_deploy_path)
        os.makedirs(self.mindie_deploy_path, mode=0o750)
        src_folder = os.path.join(self.mindie_service_path, EXAMPLES_DIR)
        if os.path.exists(src_folder):
            for item in os.listdir(src_folder):
                src_path = os.path.join(src_folder, item)
                dest_path = os.path.join(self.mindie_deploy_path, item)
                if os.path.isdir(src_path):
                    shutil.copytree(src_path, dest_path)
                else:
                    shutil.copy2(src_path, dest_path)
        else:
            self.module.fail_json(msg="[ASCEND][ERROR] source folder {} not found".format(src_folder), rc=1,
                                  changed=False)

    def run(self):
        self.unzip_package()
        self.extract_conf()

    def _unzip_run_packages(self):
        mindie_pattern = "Ascend-mindie*linux-{}*.run".format(self.arch)
        mindie_server_pattern = "Ascend-mindie-service*linux-{}*.run".format(self.arch)
        mindie_llm_pattern = "Ascend-mindie-llm*linux-{}*.run".format(self.arch)
        for pattern in [mindie_pattern, mindie_server_pattern, mindie_llm_pattern]:
            package = self.find_package(pattern)
            if not package:
                self.module.fail_json(msg="[ASCEND][ERROR] {} package not found".format(pattern), rc=1, changed=False)
            path = self.mindie_unzip_path
            if pattern == mindie_server_pattern:
                path = self.mindie_service_path
            elif pattern == mindie_llm_pattern:
                path = self.mindie_llm_path

            cmd = "bash {} --extract={}".format(package, path)
            rc, _, _ = self.module.run_command(cmd, check_rc=True)
            if rc != 0:
                self.module.fail_json(msg="[ASCEND][ERROR] extract mindie conf failed", rc=1, changed=False)

    def _unzip_whl_packages(self, llm_whl, motor_whl):
        mindie_llm_dir = os.path.join(self.mindie_unzip_path, "mindie-llm")
        mindie_motor_dir = os.path.join(self.mindie_unzip_path, "mindie-motor")

        os.makedirs(mindie_llm_dir, mode=0o750)
        os.makedirs(mindie_motor_dir, mode=0o750)

        for whl in [llm_whl, motor_whl]:
            cmd = "chmod +x {}".format(whl)
            rc, _, _ = self.module.run_command(cmd, check_rc=True)
            if rc != 0:
                self.module.fail_json(msg="[ASCEND][ERROR] chmod {} failed".format(whl), rc=1, changed=False)

        cmd = "unzip -o {} -d {}".format(motor_whl, mindie_motor_dir)
        rc, _, _ = self.module.run_command(cmd, check_rc=True)
        if rc != 0:
            self.module.fail_json(msg="[ASCEND][ERROR] unzip motor failed", rc=1, changed=False)

        cmd = "unzip -o {} -d {}".format(llm_whl, mindie_llm_dir)
        rc, _, _ = self.module.run_command(cmd, check_rc=True)
        if rc != 0:
            self.module.fail_json(msg="[ASCEND][ERROR] unzip llm failed", rc=1, changed=False)

        motor_data = self._find_data_dir(mindie_motor_dir, "mindie_motor-")
        llm_data = self._find_data_dir(mindie_llm_dir, "mindie_llm-")

        if not motor_data:
            self.module.fail_json(msg="[ASCEND][ERROR] mindie_motor .data directory not found", rc=1, changed=False)
        if not llm_data:
            self.module.fail_json(msg="[ASCEND][ERROR] mindie_llm .data directory not found", rc=1, changed=False)

        motor_purelib = os.path.join(motor_data, "purelib", "mindie_motor")
        llm_purelib = os.path.join(llm_data, "purelib", "mindie_llm")

        motor_conf = os.path.join(motor_purelib, CONF_DIR)
        motor_examples_conf = os.path.join(motor_purelib, EXAMPLES_DIR, KUBERNETES_DEPLOY_SCRIPTS, CONF_DIR)
        llm_config = os.path.join(llm_purelib, CONF_DIR, CONFIG_JSON)

        if os.path.exists(motor_conf):
            os.makedirs(motor_examples_conf, mode=0o750, exist_ok=True)
            for item in os.listdir(motor_conf):
                src = os.path.join(motor_conf, item)
                dst = os.path.join(motor_examples_conf, item)
                if os.path.isfile(src):
                    shutil.copy2(src, dst)
                elif os.path.isdir(src):
                    if os.path.exists(dst):
                        shutil.rmtree(dst)
                    shutil.copytree(src, dst)

        if os.path.exists(llm_config):
            os.makedirs(motor_examples_conf, mode=0o750, exist_ok=True)
            for name in [CONFIG_JSON, CONFIG_P_JSON, CONFIG_D_JSON]:
                shutil.copy2(llm_config, os.path.join(motor_examples_conf, name))
                os.chmod(os.path.join(motor_examples_conf, name), 0o640)

        self.mindie_service_path = motor_purelib