"""
Matilo Model Serving Worker

author: Hyungkoo.kim
"""

__author__ = "ryuvsken"
__copyright__ = "Copyright (C) 2025 Matilo (C)"


import time
import json
from datetime import datetime
import requests

from enum import Enum

from sqlalchemy import Connection, Engine, Transaction, create_engine, text

from mmsw.core.config_base import ConfigBase, set_config_base, get_config_base


"""
모델 서빙의 상태 정보

- READY : 모델 서빙 대기
- RUN : 모델 구동 중
- FAIL : 모델 구동 실패
- BADPARAM : 잘못된 매개변수
- TIMEOUT : 모델 구동 타임아웃
- CANCEL : 모델 구동 취소
- COMPLETE : 모델 구동 완료
"""
class ModelStatus(Enum):
    READY = 0
    RUN = 1
    FAIL = 2
    BADPARAM = 3
    TIMEOUT = 4
    CANCEL = 5
    COMPLETE = 200


"""
모델 워커 동작 상태

- 0: READY
- 1: RUN
- 2: TERMINATE
- 4: SHUTDOWN
"""
class WorkerStatus(Enum):
    READY = 0
    RUN = 1
    TERMINATE = 2
    SHUTDOWN = 4
    

"""
슬랙 알림 상태

- OK: 정상
- FAIL: 비정상
"""
class SlackStatus(Enum):
    OK = 0
    FAIL = 1


def run_worker(fnRunModel):

    set_config_base(ConfigBase())
    settings = get_config_base()

    slack_send(status=SlackStatus.OK, slack_msg='Worker Start: ' + ', Worker: ' + settings.MODEL_WORKER)
    worker_report(status=WorkerStatus.RUN.value, serving_id=None)

    while True:
        time.sleep(3)

        if worker_check() == False:
            break

        DB_URL = f'{settings.DB_SCHEME}://{settings.DB_USER}:{settings.DB_PWD}@{settings.DB_HOST}:{settings.DB_PORT}/{settings.DB_NAME}?charset=utf8'

        conn: Connection = None
        engine: Engine = None
        transaction: Transaction = None
        
        id: int
        model: str = None
        params: str = None

        try:
            engine = create_engine(DB_URL, echo=True)
            conn = engine.connect()
            transaction = conn.begin()
            query = f"SELECT * FROM model_serving WHERE model = '{settings.MODEL_NAME}' AND status = 0 ORDER BY id ASC LIMIT 1 FOR UPDATE;"
            rs = conn.execute(text(query)).fetchone()

            if rs is None or len(rs) == 0:
                transaction.commit()
                continue

            id=rs.id
            model=rs.model
            params=rs.params

            query = f"UPDATE model_serving SET status = {ModelStatus.RUN.value}, start_date = CURRENT_TIMESTAMP, upt_date = CURRENT_TIMESTAMP WHERE id={id};"
            conn.execute(text(query))
            transaction.commit()
        except Exception as e:
            transaction.rollback()

            slack_msg = 'Serving Check DB Query Error: ' + str(e) + ', Worker: ' + settings.MODEL_WORKER
            slack_send(status=SlackStatus.FAIL, slack_msg=slack_msg)
        finally:
            _db_safe_close(conn, engine)

        if id is not None and model is not None and params is not None:
            try:
                slack_msg = 'Start ID: ' + str(id) + ', Worker: ' + settings.MODEL_WORKER
                slack_send(status=SlackStatus.OK, slack_msg=slack_msg)

                worker_report(status=WorkerStatus.RUN.value, serving_id=id)

                status, result = fnRunModel(id=id, model_name=model, model_params=params)
                #result = result.replace('"', '\"')

                # TODO: 2. model_report() 함수를 호출하여 모델 결과 기록
                #model_report(id=id, status=200, result="{\"key1\": \"val1\"}")
                model_report(id=id, status=status.value, result=result)

                slack_msg = 'End ID: ' + str(id) + ', Worker: ' + settings.MODEL_WORKER
                slack_send(status=SlackStatus.FAIL, slack_msg=slack_msg)
            except Exception as e:
                slack_msg = 'Error ID: ' + str(id) + ", Msg: " + str(e) + ', Worker: ' + settings.MODEL_WORKER
                slack_send(status=SlackStatus.FAIL, slack_msg=slack_msg)

                model_report(id=id, status=ModelStatus.FAIL.value, result=json.dumps({"msg": str(e)}, ensure_ascii=False))


def worker_check():

    settings = get_config_base()
    DB_URL = f'{settings.DB_SCHEME}://{settings.DB_USER}:{settings.DB_PWD}@{settings.DB_HOST}:{settings.DB_PORT}/{settings.DB_NAME}?charset=utf8'

    conn: Connection = None
    engine: Engine = None

    ret = True
    try:
        engine = create_engine(DB_URL, echo=True)
        conn = engine.connect()
        query = f"SELECT * FROM model_worker WHERE name = '{settings.MODEL_WORKER}' ORDER BY id ASC LIMIT 1;"
        rs = conn.execute(text(query)).fetchone()
        if rs is None or len(rs) == 0:
            ret = False
        
        if ret and rs.status == WorkerStatus.SHUTDOWN.value:
            ret = False

    except Exception as e:
        slack_msg = 'Worker Check DB Query Error Msg: ' + str(e) + ', Worker: ' + settings.MODEL_WORKER
        slack_send(status=SlackStatus.FAIL, slack_msg=slack_msg)
    finally:
        _db_safe_close(conn, engine)

    if ret == False:
        worker_report(WorkerStatus.TERMINATE.value, None)
        
        slack_msg = 'Worker Shutdown: ' + ', Worker: ' + settings.MODEL_WORKER
        slack_send(status=SlackStatus.FAIL, slack_msg=slack_msg)


    return ret


def worker_report(status: int, serving_id: int = None):

    settings = get_config_base()
    DB_URL = f'{settings.DB_SCHEME}://{settings.DB_USER}:{settings.DB_PWD}@{settings.DB_HOST}:{settings.DB_PORT}/{settings.DB_NAME}?charset=utf8'

    conn: Connection = None
    engine: Engine = None

    try:
        id = None

        engine = create_engine(DB_URL, echo=True)
        conn = engine.connect()
        query = f"SELECT * FROM model_worker WHERE name = '{settings.MODEL_WORKER}' ORDER BY id ASC LIMIT 1;"
        rs = conn.execute(text(query)).fetchone()
        if rs is None or len(rs) == 0:
            query2 = f"INSERT INTO model_worker (name, model_name, status) VALUES ('{settings.MODEL_WORKER}', '{settings.MODEL_NAME}', 0)"
            conn.execute(text(query2))
            conn.commit()

            query = f"SELECT * FROM model_worker WHERE name = '{settings.MODEL_WORKER}' ORDER BY id ASC LIMIT 1;"
            rs = conn.execute(text(query)).fetchone()
            if rs is None or len(rs) == 0:
                raise Exception('Model Worker DB Not Exist...')
            id = rs.id
        else:
            id = rs.id

        if serving_id is not None:
            query = f"UPDATE model_worker SET status = {status}, upt_date = CURRENT_TIMESTAMP, serving_id = {serving_id}, serving_date = CURRENT_TIMESTAMP WHERE id = {id}"
        else:
            query = f"UPDATE model_worker SET status = {status}, upt_date = CURRENT_TIMESTAMP WHERE id = {id}"
        conn.execute(text(query))
        conn.commit()
    except Exception as e:
        slack_msg = 'Worker Report DB Error Msg: ' + str(e) + ', Worker: ' + settings.MODEL_WORKER
        slack_send(status=SlackStatus.FAIL, slack_msg=slack_msg)
    finally:
        _db_safe_close(conn, engine)


def model_report(id: int, status: int, result: str):

    settings = get_config_base()
    DB_URL = f'{settings.DB_SCHEME}://{settings.DB_USER}:{settings.DB_PWD}@{settings.DB_HOST}:{settings.DB_PORT}/{settings.DB_NAME}?charset=utf8'

    conn: Connection = None
    engine: Engine = None

    try:
        engine = create_engine(DB_URL, echo=True)
        conn = engine.connect()
        query = f"UPDATE model_serving SET result_msg = '{result}', status = {status}, start_date = CURRENT_TIMESTAMP, upt_date = CURRENT_TIMESTAMP WHERE id = {id};"
        conn.execute(text(query))
        conn.commit()
    except Exception as e:
        slack_msg = 'Model Report DB Error Msg: ' + str(e) + ', Worker: ' + settings.MODEL_WORKER
        slack_send(status=SlackStatus.FAIL, slack_msg=slack_msg)
    finally:
        _db_safe_close(conn, engine)
    

def _db_safe_close(conn, engine):

    try:
        if conn is not None:
            conn.close()
    except:
        pass

    try:
        if engine is not None:
            engine.dispose()
    except:
        pass


def slack_send(status: SlackStatus, slack_msg: str):

    settings = get_config_base()
    if settings.SLACK_URL is None:
        return

    try:
        url = settings.SLACK_URL
        if url is None:
            return
        
        header = {'Content-type': 'application/json'}
        username = "model-worker-bot"

        # https://slackmojis.com/
        # https://github.com/d0x2f/slack-emoji-text
        icon_emoji = ":smile:" if status == SlackStatus.OK else ":cry:"
        color = ":good:" if status == SlackStatus.OK else ":bad:"

        time_with_milliseconds_string = datetime.now().strftime("%H:%M:%S")
        # 😢
        icon_msg = ": :smile: - " if status == SlackStatus.OK else ": :cry: - "
        slack_msg = time_with_milliseconds_string + icon_msg + slack_msg
        attachments = [{
            "color": color,
            "text": slack_msg
        }]
        
        data = {"username": username, "attachments": attachments, "icon_emoji": icon_emoji}
        #print(data)

        # 메세지 전송
        requests.post(url, headers=header, json=data)
    except Exception as e:
        pass
 
 