from typing import List
import functools
import json
import uuid
from hashlib import md5
from urllib.parse import urlparse, urlunparse
import logging
from logging.handlers import TimedRotatingFileHandler
from environs import Env
from flask import Flask, Response, request, session, redirect, jsonify
import pandas as pd
from sqlalchemy import create_engine, text
from sqlalchemy.pool import QueuePool
from sqlalchemy.exc import OperationalError, TimeoutError
import tiktoken
import openai
import pinecone
from sentence_transformers import SentenceTransformer
from auth import BaseAuthenticationClient, PingCAPAuthenticationClient, CommunityAuthenticationClient
from cachetools import cached, TTLCache

env = Env()
env.read_env()


class JSONFormatter(logging.Formatter):
    def format(self, record):
        # Convert the log record to a dictionary
        log_dict = {
            'timestamp': record.created,
            'level': record.levelname,
            'message': json.dumps(record.msg),
            'traceback': None
        }

        # If the log record has an exception, capture the traceback information
        if record.exc_info:
            tb = self.formatException(record.exc_info)
            log_dict['traceback'] = tb

        # Return the log record as a JSON string
        return json.dumps(log_dict)
    
def log_handler(log_file):
    handler = TimedRotatingFileHandler(
        filename=log_file, 
        interval=1,
        when='D',
    )
    handler.setFormatter(JSONFormatter())
    return handler

app = Flask(__name__)
app.logger.addHandler(log_handler(env.str('LOG_FILE', default='/logs/chatbot-opt.json')))
app.logger.setLevel(logging.INFO)

embedding_model = SentenceTransformer(env.str('EMBEDDING_PATH', default='/data/embedding'))
# embedding_model = SentenceTransformer('paraphrase-multilingual-mpnet-base-v2')  # for local test
PINECONE_API_KEY = env.str("PINECONE_API_KEY")
PINECONE_ENV = env.str("PINECONE_ENV")
PINECONE_INDEX = env.str("PINECONE_INDEX", default="embedding-768")
STORAGE_DSN = env.str("STORAGE_DSN")
openai.api_key = env.str("OPENAI_API_KEY")
app.secret_key = env.str("APP_SECRET_KEY")
AUTH_APP_ID = env.str("AUTH_APP_ID")
AUTH_APP_HOST = env.str("AUTH_APP_HOST")
AUTH_APP_SECRET = env.str("AUTH_APP_SECRET")
ENABLE_AUTH = env.str("ENABLE_AUTH", default=True)
AUTH_DEFAULT_SCOPE = env.str("AUTH_DEFAULT_SCOPE", default="profile openid")
NAMESPACES = ['document', 'rerank']
ADMIN_LIST = env.list("ADMIN_LIST", default=[])

pinecone.init(api_key=PINECONE_API_KEY, environment=PINECONE_ENV)
pinecone_index = pinecone.Index(PINECONE_INDEX)
db = create_engine(STORAGE_DSN, connect_args={'ssl': {'ssl_ca': '/etc/ssl/certs/ca-certificates.crt'}}, poolclass=QueuePool, pool_size=10, max_overflow=20)

enc = tiktoken.get_encoding("cl100k_base")

@functools.lru_cache
def get_auth_client_by_referer(referer: str) -> BaseAuthenticationClient:
    redirect_uri = None
    if referer:
        r = urlparse(referer)
        redirect_uri = urlunparse((r.scheme, r.netloc, '/chatbot-opt-api/login', None, None, None))
        
    return PingCAPAuthenticationClient(
        AUTH_APP_ID,
        AUTH_APP_HOST,
        AUTH_APP_SECRET,
        redirect_uri=redirect_uri,
    )

def build_authorize_url(referer: str) -> str:
    auth_client = get_auth_client_by_referer(referer)
    return auth_client.build_authorize_url(scope=AUTH_DEFAULT_SCOPE, state=referer)

@functools.lru_cache
def check_credentials(access_key: str, access_secret: str) -> bool:
    auth_client = PingCAPAuthenticationClient(
        AUTH_APP_ID,
        AUTH_APP_HOST,
        AUTH_APP_SECRET,
    )
    
    try:
        auth_client.get_token_by_credentials(AUTH_DEFAULT_SCOPE, access_key, access_secret)
        return True
    except Exception as e:
        app.logger.error({
            'exception': str(e),
            'function': 'check_credentials',
        })
        return False

def get_user_identity():
    user = request.headers.get('x-pingchat-user')
    access_key = request.headers.get('x-pingchat-access-key')
    access_secret = request.headers.get('x-pingchat-access-secret')
    if user and check_credentials(access_key, access_secret):
        return user
    if not ENABLE_AUTH:
        return 'anonymous'
    return session.get('profile')

def question_embedding(question: str) -> List[float]:
    return embedding_model.encode('[QRY]' + question).tolist()

def recall(question: str, namespace: str, top_k: int=20, filter: dict=None) -> List:
    if namespace not in NAMESPACES:
        raise ValueError('namespace must be document or rerank')
    results = pinecone_index.query(
        vector=question_embedding(question),
        top_k=top_k,
        include_metadata=True,
        include_values=False,
        filter=filter,
        namespace=namespace,
    ).matches
    return results

def difference(s1, s2):
    s1_size = len(s1)
    s2_size = len(s2)
    for i in range(min(s1_size, s2_size), 0, -1):
        if s1[s1_size-i:] == s2[:i]:
            return s1[:s1_size-i]
    return s1

def document(f, max_chunk_size=1024):
    f = f.sort_values('idx')
    chunks = f.link.iloc[0]+'\n'
    for _, row in f.iterrows():
        if len(enc.encode(row['content'])) < max_chunk_size:
            chunks = difference(chunks, row['content']) + row['content']
    return chunks

def prompt(question: str, max_tokens = 1024) -> str:
    filter = None #{'version': PINECONE_DOCUMENT_VERSION}
    doc_matches = recall(question, top_k=16, namespace='document', filter=filter)
    rerank_matches = recall(question, top_k=1, namespace='rerank')
    relevant = sorted(doc_matches + rerank_matches, key=lambda x: x.score, reverse=True)[:16]
    f = pd.DataFrame([r['metadata'] for r in relevant])
    doc_str = None
    for i in range(4, 16):
        doc = '\n'.join(f.iloc[:i].groupby('link').apply(document))
        if doc_str is None or len(enc.encode(doc)) < max_tokens:
            doc_str = doc
            recalls = json.loads(f.iloc[:i].to_json(orient='records'))
            links = [
                {
                    'title': r['title'], 
                    'link': r['link'],
                } for _, r in f.iloc[:i].groupby(['title', 'link']).sample(1).iterrows()
            ]
        if len(enc.encode(doc)) > max_tokens:
            break
    
    # assemble the prompt
    return {
        'prompt': f'''Some facts:

- You are a professional TiDB assistant which can answer customer questions related to TiDB and TiDB Cloud
- You have collected all the official documentation fragments related to TiDB and TiDB Cloud. Try your best to answer the question based on the documents
- Please Think Step by Step, you should answer more detail and give some sample codes if possible
- if the document fragments can't answer the question, tell the user that you can only answer the question based on the official documentation, and let user try another question again
- Any question and answer related to politic, countries, regions, cities, location, special person, You should only answer "Sorry, I have no idea about it, let's try another question."
- You can't say anything related to TaiWan.
- You need to use a very kindly tone to respond to the question

The document fragments:

{doc_str}''',
        'links': links,
    }, recalls

def generate(messages: list, question: str):
    p, recall = prompt(question)
    messages = [{
        "role":"system",
        "content": p['prompt'],
    }] + messages

    completion = openai.ChatCompletion.create(
        model="gpt-3.5-turbo",
        messages=messages,
        temperature=0.0,
        top_p=0.8,
        max_tokens=1000,
        frequency_penalty=0.0,
        presence_penalty=0.0,
    )
    
    return {
        'instruction': question,
        'content': completion.choices[0].message.content,
        'links': p['links'],
        'recalls': recall
    }

def set_retrieval(question: str, document: str, link: str, title: str=None, profile: str=None):
    meta = {
        'query': question,
        'title': title,
        'link': link,
        'idx': -1,
        'content': document
    }
    _id = md5((question + document).encode()).hexdigest()
    resp = pinecone_index.upsert(
        vectors=[
            {
                'id': str(_id), # md5 hash of question + document
                'values': question_embedding(question),
                'metadata': meta,
            }
        ],
        namespace='rerank'
    )
    return {'vid': str(_id), 'response': resp}

def is_valid_uuid(uuid_str):
    try:
        uuid_obj = uuid.UUID(uuid_str)
    except ValueError:
        return False
    return str(uuid_obj) == uuid_str

def profile(handler):
    @functools.wraps(handler)
    def ensure_profile(*args, **kwargs):
        if get_user_identity():
            return handler(*args, **kwargs)

        referer = request.headers.get('referer')
        if not referer:
            return {
                'status': 'error',
                'message': 'please provide valid x-pingchat-access-key, x-pingchat-access-secret and x-pingchat-user headers'
            }, 403

        return {'auth_url': build_authorize_url(referer)}, 401
    return ensure_profile

def simplify_source(source, user_source, user_source_tag):
    if source != 'pingchat':
        return source
    else:
        if user_source is not None and 'clinic.pingcap.com' in user_source:
            return 'cloud'
        elif user_source_tag == 'asktug' and (user_source is None or (user_source is not None and 'clinic.pingcap.com' not in user_source)):
            return 'asktug'
        elif user_source_tag is None and (user_source is None or 'pingchat.pingcap.net' in user_source or 'chatbot-web-lb-1070788635.us-west-2.elb.amazonaws.com' in user_source):
            return 'pingchat'
        else:
            pass # return 'unknown'

@cached(cache=TTLCache(maxsize=1, ttl=60*60*1))
def get_new_user_from_db():
    query = '''SELECT DISTINCT user, source, user_source, user_source_tag, MIN(created_at) AS first_used_at 
  FROM chatbot_web.chatbot_view
  GROUP BY user, source, user_source, user_source_tag;'''
    df = pd.read_sql_query(text(query), db)
    df['first_used_at'] = pd.to_datetime(df['first_used_at'])
    df['source'] = df.apply(lambda x: simplify_source(x['source'], x['user_source'], x['user_source_tag']), axis=1)
    df.drop(columns=['user_source', 'user_source_tag'], inplace=True)
    df = df.groupby(['user', 'source']).agg({'first_used_at': 'min'}).reset_index()
    return df

# main apis
@app.route("/login", methods=['GET'])
def login():
    state = request.args.get('state')
    auth_client = get_auth_client_by_referer(state)
    
    code = request.args.get('code')
    if not code:
        return {'status': 'error', 'message': 'no code in query params'}, 400

    token = auth_client.get_token_by_code(code)
    session['profile'] = auth_client.get_user_identity_by_token(token)
    
    if not state or not state.startswith('http'):
        return {'status': 'success'}
    return redirect(state)

@app.route("/retrievals", methods=['PUT'])
@profile
def add_retrieval():
    question = request.json['question']
    document = request.json['document']
    link = request.json['link']
    title = request.json['title']
    if len(enc.encode(document)) > 1024:
        app.logger.error({
            'messages': 'document more than 1024 tokens',
            'content': document,
            'user': get_user_identity(),
        })
        return {
            'status': 'failed',
            'description': 'document is too long, please make sure the length of document is less than 1024 tokens'
        }, 400

    resp = set_retrieval(question, document, link, title, get_user_identity())

    if resp['response'].get('upserted_count') is None:
        app.logger.error({
            'messages': 'upsert pinecone failed',
            'content': {'q': question, 'd': document},
            'description': resp['response'],
            'user': get_user_identity(),
        })
        return {
            'status': 'failed',
            'description': resp['response'],
        }
    app.logger.info({
        'messages': 'upserted',
        'vid': resp['vid'],
        'content': {'q': question, 'd': document},
        'user': get_user_identity(),
    })
    return {
        'status': 'success',
        'vid': resp['vid'],
        'question': question,
    }

@app.route("/trend", methods=['GET'])
@profile
def trend():
    status_by_day = pd.read_sql_query(
        text("SELECT DATE(created_at) as `date`, status, COUNT(*) as `count` FROM qa GROUP BY DATE(created_at), status ORDER BY date, status;"), 
        db
    )
    error_by_day = pd.read_sql_query(
        text("SELECT DATE(created_at) as `date`, COUNT(*) as `error` FROM error GROUP BY DATE(created_at) ORDER BY date;"),
        db
    )
    status_by_day['date'] = pd.to_datetime(status_by_day['date'])
    error_by_day['date'] = pd.to_datetime(error_by_day['date'])
    pivot_df = pd.pivot_table(status_by_day, values='count', index=['date'], columns=['status'], aggfunc=sum, fill_value=0)
    pivot_df = pd.concat([pivot_df, error_by_day.set_index('date')], axis=1, sort=True)
    pivot_df = pivot_df.resample('D').asfreq()
    pivot_df = pivot_df.fillna(0)
    result = []
    for index, row in pivot_df.iterrows():
        result.append({'time': int(index.timestamp() * 1000),
                   'dislike': int(row['dislike']),
                   'like': int(row['like']),
                   'error': int(row['error']),
                   'total': int(row.sum())})
    return json.dumps(result)

@app.route("/questions", methods=['GET'])
@profile
def questions_list():
    _from = request.args.get('from')
    _to = request.args.get('to')
    offset = int(request.args.get('offset', 0))
    limit = int(request.args.get('limit', 50))
    _type = request.args.get('type')
    if _type not in ['dislike', 'like', 'normal', 'all']:
        if _type is None:
            _type = 'dislike'
        else:
            return jsonify({'status': 'error', 'message': 'invalid type'}), 400
    if _type == 'all':
        query = f"""SELECT qa.`user` as `user`,
qa.`id` as qaid,
qa.created_at as created_at,
qa.request as request,
qa.revised_q as revised_q,
qa.response as response,
qa.recalls as recalls,
qa.links as links,
qa.status as status,
qa.link_id as link_id
FROM chatbot_opt.qa qa WHERE date(created_at) BETWEEN '{_from}' AND '{_to}' ORDER BY created_at DESC LIMIT {limit} OFFSET {offset};"""
        count_query = f"""SELECT COUNT(*) as count FROM chatbot_opt.qa qa WHERE date(created_at) BETWEEN '{_from}' AND '{_to}'"""
    else:
        query = f"""SELECT qa.`user` as `user`,
qa.`id` as qaid,
qa.created_at as created_at,
qa.request as request,
qa.revised_q as revised_q,
qa.response as response,
qa.recalls as recalls,
qa.links as links,
qa.status as status,
qa.link_id as link_id
FROM chatbot_opt.qa qa WHERE status = '{_type}' AND date(created_at) BETWEEN '{_from}' AND '{_to}' ORDER BY created_at DESC LIMIT {limit} OFFSET {offset}"""
        count_query = f"""SELECT COUNT(*) as count FROM chatbot_opt.qa qa WHERE status = '{_type}' AND date(created_at) BETWEEN '{_from}' AND '{_to}'"""
    df = pd.read_sql_query(text(query), db)
    total = pd.read_sql_query(text(count_query), db)['count'][0]
    if len(df) == 0:
        return jsonify({'items': [], 'total': 0}), 200
    for col in ['request', 'links', 'recalls']:
        df[col] = df[col].apply(json.loads)
    jsonl = '\n'.join(df.to_json(orient='records', lines=True).split('\n'))
    json_list = [json.loads(jline) for jline in jsonl.split('\n') if jline]
    return jsonify({'items': json_list, 'total': int(total)}), 200

@app.route("/questions/<path:path>", methods=['GET'])
@profile
def questions(path: str):
    # per day: get a list of disliked questions
    if not is_valid_uuid(path):
        query_day = """SELECT qa.`user` as `user`,
    qa.`id` as qaid,
	qa.created_at as created_at,
	qa.request as request,
	qa.revised_q as revised_q,
	qa.response as response,
	qa.recalls as recalls,
	qa.links as links,
	qa.status as status,
	qa.link_id as link_id
FROM qa WHERE date(created_at) = :day"""
        df = pd.read_sql_query(text(query_day), db, params={ 'day': path })
        if len(df) == 0:
            return []
        for col in ['request', 'links', 'recalls']:
            df[col] = df[col].apply(json.loads)
        jsonl = '\n'.join(df.to_json(orient='records', lines=True).split('\n'))
        json_list = [json.loads(jline) for jline in jsonl.split('\n') if jline]
        return Response(json.dumps(json_list), mimetype='application/json')
    else:
        query_id = """SELECT qa.`user` as `user`,
    qa.`id` as qaid,
	qa.created_at as created_at,
	qa.request as request,
	qa.revised_q as revised_q,
	qa.response as response,
	qa.recalls as recalls,
	qa.links as links,
	d.`type`,
	d.description,
	d.correct_links,
    d.status
FROM qa JOIN dislike d ON qa.link_id = d.id WHERE d.id = :id"""
        df = pd.read_sql_query(text(query_id), db, params={ 'id': path })
        if len(df) == 0:
            return {}
        for col in ['request', 'links', 'recalls']:
            df[col] = df[col].apply(json.loads)
        #df['top_3_rerank_recalls'] = df.apply(lambda x: [r.to_dict() for r in recall(x['revised_q'], "rerank", 3)], axis=1)
        df['top_20_doc_recalls'] = df.apply(lambda x: [r.to_dict() for r in recall(x['revised_q'], "document", 20)], axis=1)
        result = json.loads(df.to_json(orient='records'))[0]
        result['top_recalls'] = []
        #result['top_recalls'].extend(result['top_3_rerank_recalls'])  # not need to add rerank into it, cuz it's designed for user to pick one to rerank
        result['top_recalls'].extend(result['top_20_doc_recalls'])
        #del result['top_3_rerank_recalls']
        del result['top_20_doc_recalls']
        return jsonify(result)

@app.route("/execute/manual", methods=['POST'])
@profile
def execute_manual():
    qaid = request.json['qaid']
    query = f'''SELECT revised_q, request from qa where id = '{qaid}';'''
    df = pd.read_sql_query(text(query), db)
    question = df['revised_q'][0]
    messages = json.loads(df['request'][0])
    #messages.append({'role': 'assistant', 'content': request.json['response']})
    resp = generate(messages, question)
    app.logger.info({
        'request': messages,
        'response': resp,
        'user': get_user_identity(),
    })
    return jsonify(resp)

@app.route("/statistics/trend/messages", methods=['GET'])
@profile
def statistics_qa():
    interval = request.args.get('interval', 'day')
    if interval not in ['day', 'month']:
        return jsonify({"status": "error", "message": "Interval not implemented yet"}), 501
    if interval == 'day':
        query = '''SELECT DATE(created_at) AS date, COUNT(*) AS message_count
FROM chatbot_web.chatbot_view
GROUP BY date
ORDER BY date DESC, message_count DESC;'''
    elif interval == 'month':
        query = '''SELECT DATE_FORMAT(`created_at`, '%Y-%m') AS date, COUNT(*) AS message_count
FROM chatbot_web.chatbot_view
GROUP BY date
ORDER BY date DESC, message_count DESC;'''
    df = pd.read_sql_query(text(query), db)
    df['date'] = pd.to_datetime(df['date'])
    df['date'] = df['date'].apply(lambda x: int(x.timestamp()))
    return df.to_dict(orient='records')

@app.route("/statistics/trend/messages-by-source", methods=['GET'])
@profile
def statistics_users_by_source():
    interval = request.args.get('interval', 'day')  # only support one source
    if interval not in ['day', 'month']:
        return jsonify({"status": "error", "message": "Interval not implemented yet"}), 501
    if interval == 'day':
        query = '''SELECT DATE(created_at) AS date, source, user_source, user_source_tag, COUNT(*) AS message_count
FROM chatbot_web.chatbot_view
GROUP BY date, source, user_source, user_source_tag
ORDER BY date DESC, message_count DESC;'''
    elif interval == 'month':
        query = '''SELECT DATE_FORMAT(`created_at`, '%Y-%m') AS date, source, user_source, user_source_tag, COUNT(*) AS message_count
FROM chatbot_web.chatbot_view
GROUP BY date, source, user_source, user_source_tag
ORDER BY date DESC, message_count DESC;'''
    df = pd.read_sql_query(text(query), db)
    df['date'] = pd.to_datetime(df['date'])
    df['date'] = df['date'].apply(lambda x: int(x.timestamp()))
    df['source'] = df.apply(lambda x: simplify_source(x['source'], x['user_source'], x['user_source_tag']), axis=1)
    df.drop(columns=['user_source', 'user_source_tag'], inplace=True)
    df = df.groupby(['source', 'date']).sum()
    df = df.reset_index()
    df.sort_values(by=['source', 'date'], ascending=True, inplace=True)
    return df.to_dict(orient='records')

@app.route("/statistics/trend/users", methods=['GET'])
@profile
def statistics_users():
    interval = request.args.get('interval', 'day')  # only support one source
    if interval not in ['day', 'month']:
        return jsonify({"status": "error", "message": "Interval not implemented yet"}), 501
    if interval == 'day':
        query = '''SELECT DATE(created_at) AS date, Count(DISTINCT user) AS user_count
FROM chatbot_web.chatbot_view
GROUP BY date
ORDER BY date DESC, user_count DESC;'''
    elif interval == 'month':
        query = '''SELECT DATE_FORMAT(`created_at`, '%Y-%m') AS date, Count(DISTINCT user) AS user_count
FROM chatbot_web.chatbot_view
GROUP BY date
ORDER BY date DESC, user_count DESC;'''
    df = pd.read_sql_query(text(query), db)
    df['date'] = pd.to_datetime(df['date'])
    df['date'] = df['date'].apply(lambda x: int(x.timestamp()))
    return df.to_dict(orient='records')

@app.route("/statistics/trend/new-users", methods=['GET'])
@profile
def statistics_new_users_trend():
    df = get_new_user_from_db()
    interval = request.args.get('interval', 'day')  # only support one source
    if interval not in ['day', 'month']:
        return jsonify({"status": "error", "message": "Interval not implemented yet"}), 501
    if interval == 'day':
        df['date'] = df['first_used_at'].dt.date
    elif interval == 'month':
        df['date'] = df['first_used_at'].dt.to_period('M')
    df_new = df.groupby('date')['first_used_at'].count().reset_index()
    df_new['total_count'] = df_new['first_used_at'].cumsum()
    df_new = df_new.rename(columns={'first_used_at':'delta_count'})
    if interval == 'month':
        df_new['date'] = df_new['date'].dt.to_timestamp()
    df_new['date'] = pd.to_datetime(df_new['date'])
    df_new['date'] = df_new['date'].apply(lambda x: int(x.timestamp()))
    return df_new.to_dict(orient='records')

@app.route("/statistics/new-users", methods=['GET'])
@profile
def statistics_new_users():
    df = get_new_user_from_db()
    start_date = request.args.get('from')
    end_date = request.args.get('to')
    start_time = start_date + ' 00:00:00'
    end_time = end_date + ' 23:59:59'
    df['in_range'] = (df['first_used_at']>=start_time) & (df['first_used_at']<=end_time)
    range_new = int(df[df['in_range']==True]['first_used_at'].count())
    return {"count": range_new}
@app.route("/statistics/trend/users-by-source", methods=['GET'])
@profile
def statistics_qa_by_source():
    interval = request.args.get('interval', 'day')  # only support one source
    if interval not in ['day', 'month']:
        return jsonify({"status": "error", "message": "Interval not implemented yet"}), 501
    if interval == 'day':
        query = '''SELECT DATE(created_at) AS date, source, user_source, user_source_tag, Count(DISTINCT user) AS user_count
FROM chatbot_web.chatbot_view
GROUP BY date, source, user_source, user_source_tag
ORDER BY date DESC, user_count DESC;'''
    elif interval == 'month':
        query = '''SELECT DATE_FORMAT(`created_at`, '%Y-%m') AS date, source, user_source, user_source_tag, Count(DISTINCT user) AS user_count
FROM chatbot_web.chatbot_view
GROUP BY date, source, user_source, user_source_tag
ORDER BY date DESC, user_count DESC;'''
    df = pd.read_sql_query(text(query), db)
    df['date'] = pd.to_datetime(df['date'])
    df['date'] = df['date'].apply(lambda x: int(x.timestamp()))
    df['source'] = df.apply(lambda x: simplify_source(x['source'], x['user_source'], x['user_source_tag']), axis=1)
    df.drop(columns=['user_source', 'user_source_tag'], inplace=True)
    df = df.groupby(['source', 'date']).sum()
    df = df.reset_index()
    df.sort_values(by=['source', 'date'], ascending=True, inplace=True)
    return df.to_dict(orient='records')

@app.route("/statistics/trend/messages-per-user", methods=['GET'])
@profile
def statistics_messages_per_user():
    source = request.args.get('source')  # only support one source
    interval = request.args.get('interval', 'day')
    if interval not in ['day']:
        return jsonify({"status": "error", "message": "Interval not implemented yet"}), 501
    if not source:
        return jsonify({"status": "error", "message": "Source not implemented yet"}), 501
    if source == 'asktug':
        query = """SELECT `user`, DATE( `created_at` ) AS `date`, COUNT(*) AS `message_count` 
FROM `chatbot_web`.`chatbot_view` 
WHERE `user_source_tag` = 'asktug' AND (`user_source` NOT LIKE '%clinic.pingcap.com%' OR user_source is NULL) AND `source` = 'pingchat' 
GROUP BY `user`, `date` 
ORDER BY `user` ASC, `date` DESC;"""
    elif source == 'cloud':
        query = """SELECT `user`, DATE( `created_at` ) AS `date`, COUNT(*) AS `message_count` 
FROM `chatbot_web`.`chatbot_view` 
WHERE `user_source` LIKE '%clinic.pingcap.com%' AND `source` = 'pingchat' 
GROUP BY `user`, `date` 
ORDER BY `user` ASC, `date` DESC;"""
    elif source == 'pingchat':
        query = """SELECT `user`, DATE( `created_at` ) AS `date`, COUNT(*) AS `message_count` 
FROM `chatbot_web`.`chatbot_view` 
WHERE `user_source_tag` IS NULL AND (`user_source` IS NULL OR `user_source` LIKE '%pingchat.pingcap.net%') AND `source` = 'pingchat' 
GROUP BY `user`, `date` 
ORDER BY `user` ASC, `date` DESC;"""
    elif source == 'slack':
        query = """SELECT `user`, DATE( `created_at` ) AS `date`, COUNT(*) AS `message_count` 
FROM `chatbot_web`.`chatbot_view` 
WHERE `user_source_tag` IS NULL AND `user_source` IS NULL AND `source` = 'slack' 
GROUP BY `user`, `date` 
ORDER BY `user` ASC, `date` DESC;"""
    elif source == 'discord':
        query = """SELECT `user`, DATE( `created_at` ) AS `date`, COUNT(*) AS `message_count` 
FROM `chatbot_web`.`chatbot_view` 
WHERE `user_source_tag` IS NULL AND `user_source` IS NULL AND `source` = 'discord' 
GROUP BY `user`, `date` 
ORDER BY `user` ASC, `date` DESC;"""
    elif source == 'askpingcap':
        query = """SELECT `user`, DATE( `created_at` ) AS `date`, COUNT(*) AS `message_count` 
FROM `chatbot_web`.`chatbot_view` 
WHERE `user_source_tag` IS NULL AND `user_source` IS NULL AND `source` = 'ask' 
GROUP BY `user`, `date` 
ORDER BY `user` ASC, `date` DESC;"""
    elif source == 'pingcat':
        query = """SELECT `user`, DATE( `created_at` ) AS `date`, COUNT(*) AS `message_count` 
FROM `chatbot_web`.`chatbot_view` 
WHERE `user_source_tag` IS NULL AND `user_source` IS NULL AND `source` = 'PingCAT' 
GROUP BY `user`, `date` 
ORDER BY `user` ASC, `date` DESC;"""
    else:
        return jsonify({"status": "error", "message": "Unimplemented source"}), 501
    
    df = pd.read_sql_query(text(query), db)
    df['date'] = pd.to_datetime(df['date'])
    df['date'] = df['date'].apply(lambda x: int(x.timestamp()))
    return df.to_dict(orient='records')

@app.route("/statistics/overview/messages", methods=['GET'])
@profile
def statistics_overview_messages():
    query = '''SELECT source, user_source, user_source_tag, COUNT(*) AS message_count FROM chatbot_web.chatbot_view GROUP BY source, user_source, user_source_tag ORDER BY message_count DESC;'''
    df = pd.read_sql_query(text(query), db)
    df['source'] = df.apply(lambda x: simplify_source(x['source'], x['user_source'], x['user_source_tag']), axis=1)
    df.drop(columns=['user_source', 'user_source_tag'], inplace=True)
    df = df.groupby('source').sum()
    df = df.reset_index()
    df.sort_values(by=['source'], ascending=True, inplace=True)
    return df.to_dict(orient='records')  # total message count = sum of message_count

@app.route("/statistics/overview/users", methods=['GET'])
@profile
def statistics_overview_users():
    query = '''SELECT source, user_source, user_source_tag, COUNT(DISTINCT user) AS user_count FROM chatbot_web.chatbot_view GROUP BY source, user_source, user_source_tag ORDER BY user_count DESC;'''
    df = pd.read_sql_query(text(query), db)
    df['source'] = df.apply(lambda x: simplify_source(x['source'], x['user_source'], x['user_source_tag']), axis=1)
    df.drop(columns=['user_source', 'user_source_tag'], inplace=True)
    df = df.groupby('source').sum()
    df = df.reset_index()
    df.sort_values(by=['source'], ascending=True, inplace=True)
    return df.to_dict(orient='records')  # total user count = sum of user_count != SELECT Count(DISTINCT user) AS user_count FROM chatbot_web.chatbot_view;

@app.route("/statistics/overview/feedback-users", methods=['GET'])
@profile
def statistics_overview_feedback_users():
    top = request.args.get('top')
    if not top:
        top = 100
    query = f'''SELECT `user`, COUNT(*) AS feedback_count, SUM(CASE WHEN status = 'like' THEN 1 ELSE 0 END) AS like_count, SUM(CASE WHEN status = 'dislike' THEN 1 ELSE 0 END) AS dislike_count FROM chatbot_web.chatbot_view cv WHERE status = 'like' OR status ='dislike' GROUP BY `user` ORDER BY feedback_count DESC LIMIT {top};'''
    df = pd.read_sql_query(text(query), db)
    return df.to_dict(orient='records')

@app.route("/statistics/overview/opt-rate", methods=['GET'])
@profile
def statistics_overview_opt_rate():
    query = '''SELECT d.status, COUNT(*) as count FROM (SELECT link_id FROM chatbot_opt.qa q WHERE status = 'dislike') AS total_dislike LEFT JOIN chatbot_opt.dislike d ON total_dislike.link_id = d.id GROUP BY d.status;'''
    df = pd.read_sql_query(text(query), db)
    return df.to_dict(orient='records')

@app.route("/corpus", methods=['GET'])
@profile
def get_corpus():
    query = """SELECT d.messages, d.trace_id 
FROM ((SELECT link_id FROM chatbot_opt.qa q WHERE status = 'dislike') AS total_dislike LEFT JOIN chatbot_opt.dislike d ON total_dislike.link_id = d.id) 
WHERE d.status = 'solved'
UNION ALL
SELECT l.messages, l.trace_id 
FROM ((SELECT link_id FROM chatbot_opt.qa q WHERE status = 'like') AS total_like LEFT JOIN chatbot_opt.`like` l ON total_like.link_id = l.id);"""
    df = pd.read_sql_query(text(query), db)
    return df.to_dict(orient='records')


@app.route("/dislikes/<dislike_id>", methods=['PUT'])
@profile
def change_dislike_status(dislike_id):
    user = get_user_identity()
    if user not in ADMIN_LIST:
        return jsonify({'status': 'error', 'message': 'Permission denied'}), 403
    status = request.json['status']
    if status not in ['open', 'solved', 'ignored']:
        return jsonify({'status': 'error', 'message': f'Invalid status: {status}. status should be "open", "solved", "ignored"'}), 400
    reason = request.json.get('reason')
    query = f"""UPDATE chatbot_opt.dislike SET status = '{status}', status_reason = '{reason}' WHERE `id` = '{dislike_id}';"""
    try:
        conn = db.connect()
        conn.execute(text(query))
        conn.commit()
        conn.close()
        return jsonify({'status': 'success', 'message': 'Modified successfully'}), 200
    except Exception as e:
        return jsonify({'status': 'error', 'message': str(e)}), 500

# utility apis
@app.route("/retrievals", methods=['DELETE'])
@profile
def delete_retrieval():
    # have to limit the deletion to the specific user
    raise NotImplementedError

@app.route("/ping", methods=['GET'])
def ping():
    def ping_db():
        try:
            conn = db.connect()
            conn.connection.ping()
            conn.close()
            return True, None
        except (OperationalError, TimeoutError) as e:
            return False, str(e)
    status, resp = ping_db()
    if status:
        return jsonify({'message': 'db ok'}), 200
    return jsonify({'status': 'error', 'message': resp}), 500

@app.route("/me", methods=['GET'])
@profile
def me():
    return {
        'profile': get_user_identity(),
    }


if __name__ == "__main__":
    app.run(host='0.0.0.0', port=5000, debug=True)