import uuid
import logging
import json
import base64
import hashlib
import string
import random
import requests
import functools
import datetime
from environs import Env
from urllib.parse import urlparse, urlunparse
from flask import Flask, session, request, redirect, jsonify, Response
from flask_cors import CORS, cross_origin
from logs import log_handler

from auth import (
    BaseAuthenticationClient, 
    PingCAPAuthenticationClient, 
    CommunityAuthenticationClient,
)

from models import (
    query_sessions,
    count_sessions,
    insert_session,
    check_session_owner,
    query_session,
    update_session,
    is_session_valid_public,
    is_interaction_belongs_to_session,
    query_interactions,
    insert_interaction,
    update_interaction,
    query_interaction,
    find_latest_interaction,
    add_share,
    delete_share,
    query_share,
    insert_feedback,
    delete_feedback,
    ping_db,
)


env = Env()
env.read_env()

CHATBOT_SERVICE_ENDPOINT = env.str("CHATBOT_SERVICE_ENDPOINT")
STORAGE_DSN = env.str("STORAGE_DSN")
APP_SECRET_KEY = env.str("APP_SECRET_KEY")
PINGCHAT_ACCESS_KEY = env.str("PINGCHAT_ACCESS_KEY")
PINGCHAT_ACCESS_SECRET = env.str("PINGCHAT_ACCESS_SECRET")
COMMUNITY_FRONTEND_HOST = env.str("COMMUNITY_FRONTEND_HOST")
COMMUNITY_AUTH_APP_HOST = env.str("COMMUNITY_AUTH_APP_HOST")
COMMUNITY_AUTH_APP_ID = env.str("COMMUNITY_AUTH_APP_ID")
COMMUNITY_AUTH_APP_SECRET = env.str("COMMUNITY_AUTH_APP_SECRET")
AUTH_APP_ID = env.str("AUTH_APP_ID")
AUTH_APP_HOST = env.str("AUTH_APP_HOST")
AUTH_APP_SECRET = env.str("AUTH_APP_SECRET")
LINGUFLOW_ENDPOINT = env.str("LINGUFLOW_ENDPOINT")
LANGLINK_ACCESS_KEY = env.str("LANGLINK_ACCESS_KEY")
LANGLINK_ACCESS_SECRET = env.str("LANGLINK_ACCESS_SECRET")

app = Flask(__name__)
CORS(app, resources={r"/*": {"origins": ["http://localhost:3000"], "methods": ["POST", "PUT", "OPTIONS", "GET", "DELETE", "PATCH"], "allow_headers": "*"}}, supports_credentials=True)
app.config['SECRET_KEY'] = APP_SECRET_KEY
app.logger.addHandler(log_handler(env.str('LOG_FILE', default='/data/logs/chatbot-web.json')))
app.logger.setLevel(logging.INFO)


@app.after_request
def add_cors_headers(response):
    response.headers['Access-Control-Allow-Methods'] = 'POST,PUT,PATCH,GET,DELETE,OPTIONS'
    response.headers['Access-Control-Allow-Headers'] = '*'
    response.headers['Access-Control-Allow-Credentials'] = 'true'
    response.headers['Access-Control-Allow-Origin'] = "*"
    return response

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

@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-api/login', None, None, None))
        
        if r.netloc == COMMUNITY_FRONTEND_HOST:
            return CommunityAuthenticationClient(
                COMMUNITY_AUTH_APP_ID,
                COMMUNITY_AUTH_APP_HOST,
                COMMUNITY_AUTH_APP_SECRET,
                redirect_uri=redirect_uri,
            )

    return PingCAPAuthenticationClient(
        AUTH_APP_ID,
        AUTH_APP_HOST,
        AUTH_APP_SECRET,
        redirect_uri=redirect_uri,
    )

@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("profile openid", access_key, access_secret)
        return True
    except Exception as e:
        app.logger.error({
            'exception': str(e),
            'function': 'check_credentials',
        })
        return False

@functools.lru_cache(maxsize=16)
def get_clinic_user(clinic_cookie: str, csrf_token: str) -> str:
    profile_url = "https://clinic.pingcap.com/clinic/api/v1/user/profile"
    user = requests.get(profile_url, headers={
        'cookie': clinic_cookie,
        'x-csrf-token': csrf_token,
    }).json()
    return user.get('email')

def get_user_identity():
    cloud_user_info = request.headers.get('x-cloud-user')
    access_key = request.headers.get('x-pingchat-access-key')
    access_secret = request.headers.get('x-pingchat-access-secret')
    if cloud_user_info and check_credentials(access_key, access_secret):
        user = json.loads(base64.b64decode(cloud_user_info).decode())
        return user['email']

    clinic_cookie = request.headers.get('x-clinic-cookie')
    csrf_token = request.headers.get('x-csrf-token')
    if clinic_cookie:
        return get_clinic_user(clinic_cookie, csrf_token)

    return session.get('profile')

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


def generate_password(length):
    characters = string.ascii_lowercase + string.digits
    password = ''.join(random.choice(characters) for _ in range(length))
    return password


def calculate_md5(string):
    md5_hash = hashlib.md5()
    md5_hash.update(string.encode('utf-8'))
    return md5_hash.hexdigest()

def get_parent_nodes(interactions, interaction_id, depth=None):
    if not isinstance(depth, int):
        depth = len(interactions)
    
    result = []
    current_id = interaction_id

    while current_id is not None and depth > 0:
        if current_id in interactions:
            parent_id = interactions[current_id]["parent_id"]
            if parent_id is not None and parent_id in interactions:
                parent_node = interactions[parent_id]
                result.append(parent_node)
                current_id = parent_id
                if depth > 1:
                    depth -= 1
                else:
                    break
            else:
                current_id = None
        else:
            current_id = None
    result.reverse()
    return result

def get_historical_msgs(interactions, interaction_id, depth=None):
    result = []
    nodes = get_parent_nodes(interactions, interaction_id, depth)
    for node in nodes:
        result.append({"content": node["request"], "role": "user"})
        result.append({"content": node["response"], "role": "assistant"})
    return result

@app.errorhandler(Exception)
def handle_error(error):
    app.logger.error({'message': str(error), 'function': request.endpoint})
    return jsonify({'status': 'error', 'message': str(error)}), 500


@app.errorhandler(404)
def handle_not_found_error(error):
    return jsonify({'status': 'error', 'message': 'Not Found'}), 404


@app.route("/login", methods=['GET'])
@cross_origin(origins=["http://localhost:3000"], methods=["POST", "PUT", "OPTIONS", "GET", "DELETE", "PATCH"], all_headers="*", supports_credentials=True, automatic_options=True)
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('/me', methods=['GET'])
@cross_origin(origins=["http://localhost:3000"], methods=["POST", "PUT", "OPTIONS", "GET", "DELETE", "PATCH"], all_headers="*", supports_credentials=True, automatic_options=True)
@profile
def me():
    return {
        'profile': get_user_identity(),
    }

@app.route('/sessions', methods=['GET'])
@cross_origin(origins=["http://localhost:3000"], methods=["POST", "PUT", "OPTIONS", "GET", "DELETE", "PATCH"], all_headers="*", supports_credentials=True, automatic_options=True)
@profile
def get_sessions():
    user_hash = calculate_md5(get_user_identity() + ' ' + request.headers.get('host'))
    offset = int(request.args.get('offset', 0))
    limit = int(request.args.get('limit', 20))
    order = request.args.get('order', 'desc')

    sessions = query_sessions(user_hash, offset, limit, order)
    total = count_sessions(user_hash)

    response = {
        'items': sessions,
        'limit': limit,
        'offset': offset,
        'total': total,
    }
    return jsonify(response), 200
    

@app.route('/sessions/new', methods=['POST'])
@cross_origin(origins=["http://localhost:3000"], methods=["POST", "PUT", "OPTIONS", "GET", "DELETE", "PATCH"], all_headers="*", supports_credentials=True, automatic_options=True)
@profile
def create_session():
    user_email = get_user_identity()
    user_source = request.headers.get('host')
    user_hash = calculate_md5(user_email + ' ' + user_source)
    session_id = str(uuid.uuid4())
    created_at = datetime.datetime.utcnow()
    title = session_id

    # Insert new session into the database
    insert_session(session_id, user_hash, user_email, user_source, created_at, title)

    response = {
        'session_id': session_id
    }

    return jsonify(response)

@app.route("/sessions/<session_id>", methods=['POST'])
@cross_origin(origins=["http://localhost:3000"], methods=["POST", "PUT", "OPTIONS", "GET", "DELETE", "PATCH"], all_headers="*", supports_credentials=True, automatic_options=True)
@profile
def chat_stream(session_id):
    host = request.headers.get('host')
    user_hash = calculate_md5(get_user_identity() + ' ' + host)
    if not check_session_owner(session_id, user_hash):
        return jsonify({"status": "error", "message": "Session not found or does not belong to the current user"}), 404

    parent_id = request.json.get('parent_id')
    if parent_id is not None and (not is_interaction_belongs_to_session(session_id, parent_id)):
        return jsonify({"status": "error", "message": "Parent interaction does not belong to the session"}), 400
    messages = request.json.get('messages')
    interactions = query_interactions(session_id)
    if parent_id is not None:
        # get parent and grandparent if exists
        parent = interactions.get(parent_id)
        if parent is not None:
            if parent.get('request') is not None and parent.get('response') is not None:
                messages.insert(0, {"role": "assistant", "content": parent['response']})
                messages.insert(0, {"role": "user", "content": parent['request']})
            if parent['parent_id'] is not None:
                grandparent = interactions.get(parent['parent_id'])
                if grandparent is not None:
                    parent['parent'] = grandparent
                    if grandparent.get('request') is not None and grandparent.get('response') is not None:
                        messages.insert(0, {"role": "assistant", "content": grandparent['response']})
                        messages.insert(0, {"role": "user", "content": grandparent['request']})
    else:
        # rename the session title
        title = messages[0]['content']
        payload = {"title": title[:20]}
        update_session(session_id, payload)
    current_id = str(uuid.uuid4())
    created_at = datetime.datetime.utcnow()

    headers = {
        'x-pingchat-access-key': PINGCHAT_ACCESS_KEY,
        'x-pingchat-access-secret': PINGCHAT_ACCESS_SECRET,
        'x-pingchat-user': get_user_identity(),
        'x-trace-id': current_id,
    }

    params = {}
    namespaces = request.json.get('domain')
    if "clinic.pingcap.com" in host:
        params['domain'] = 'tidbcloud'
    if isinstance(namespaces, list) and len(namespaces) > 0:
        params['domain'] = namespaces
    version = request.json.get('version')
    if version is not None:
        params['version'] = version
    response = requests.post(f'{CHATBOT_SERVICE_ENDPOINT}/chat-stream', json=messages, headers=headers, params=params, stream=True)

    def generate():
        content = ''
        #for chunk in response.iter_content():
        for chunk in response.iter_lines():
            response_json = json.loads(chunk)
            if response_json.get('content') is not None:
                content += response_json['content']
                yield chunk.decode('utf-8') + '\n'
            else:
                recalls = response_json.get('recalls')
                revised_q = response_json.get('instruction')
                links = response_json.get('links')
                yield json.dumps({"links": links}) + '\n'

        score = 0
        error = ''
        updated_at = datetime.datetime.utcnow()
        insert_interaction(current_id, session_id, created_at, updated_at, parent_id, messages[-1]['content'], content, revised_q, recalls, links, score, error)
        yield json.dumps({"parent_id": parent_id, "current_id": current_id})
    return app.response_class(generate(), mimetype='text/plain')


@app.route("/sessions/<session_id>", methods=['GET'])
@cross_origin(origins=["http://localhost:3000"], methods=["POST", "PUT", "OPTIONS", "GET", "DELETE", "PATCH"], all_headers="*", supports_credentials=True, automatic_options=True)
@profile
def get_session(session_id):
    # Retrieve session details from the database
    user_hash = calculate_md5(get_user_identity() + ' ' + request.headers.get('host'))
    if not check_session_owner(session_id, user_hash):
        return jsonify({"status": "error", "message": "Session not found or does not belong to the current user"}), 404
        
    conversation = query_session(session_id)

    if conversation is None:
        return jsonify(error='Session not found')
    interactions = query_interactions(session_id)
    conversation['interactions'] = interactions
    latest_interaction_id, latest_interaction_updatedtime = find_latest_interaction(session_id)
    conversation['updated_at'] = latest_interaction_updatedtime
    conversation['latest_interaction'] = latest_interaction_id

    return jsonify(conversation), 200

@app.route("/sessions/<session_id>", methods=['DELETE'])
@cross_origin(origins=["http://localhost:3000"], methods=["POST", "PUT", "OPTIONS", "GET", "DELETE", "PATCH"], all_headers="*", supports_credentials=True, automatic_options=True)
@profile
def delete_session(session_id):
    user_hash = calculate_md5(get_user_identity() + ' ' + request.headers.get('host'))
    if not check_session_owner(session_id, user_hash):
        return jsonify({"status": "error", "message": "Session not found or does not belong to the current user"}), 404
    
    payload = {"is_deleted": True}
    update_session(session_id, payload)
    delete_share(session_id)
    check_session_owner.cache_clear()

    return jsonify({"message": "Session deleted"}), 200

@app.route("/sessions/<session_id>", methods=['PUT'])
@cross_origin(origins=["http://localhost:3000"], methods=["POST", "PUT", "OPTIONS", "GET", "DELETE", "PATCH"], all_headers="*", supports_credentials=True, automatic_options=True)
@profile
def update_session_title(session_id):
    # Didn't check if title contains sensitive words
    user_hash = calculate_md5(get_user_identity() + ' ' + request.headers.get('host'))
    if not check_session_owner(session_id, user_hash):
        return jsonify({"status": "error", "message": "Session not found or does not belong to the current user"}), 404

    # Get the new title from the request body
    new_title = request.json.get("title")
    payload = {"title": new_title}
    update_session(session_id, payload)

    return jsonify({"message": "Session title updated"}), 200


@app.route("/sessions/<session_id>/share", methods=['POST'])
@cross_origin(origins=["http://localhost:3000"], methods=["POST", "PUT", "OPTIONS", "GET", "DELETE", "PATCH"], all_headers="*", supports_credentials=True, automatic_options=True)
@profile
def share_session(session_id):
    user_hash = calculate_md5(get_user_identity() + ' ' + request.headers.get('host'))
    if not check_session_owner(session_id, user_hash):
        return jsonify({"status": "error", "message": "Session not found or does not belong to the current user"}), 404
    is_shared, share = is_session_valid_public(session_id)
    if is_shared:
        return jsonify(
            thread=json.loads(share['thread']),
            expired_at=share['expired_at'],
            password=share['password'],
            share_id=share['share_id'],
            status='old'
        ),200
    latest_interaction_id = request.json.get("latest_interaction")
    interactions = query_interactions(session_id)

    # Find the latest interaction and trace the thread
    thread = []
    current_interaction_id = latest_interaction_id
    while current_interaction_id:
        interaction = interactions.get(current_interaction_id)
        if interaction:
            thread.insert(0, interaction)
            current_interaction_id = interaction["parent_id"]
        else:
            break
    # Convert the thread to JSON
    thread_json = json.dumps(thread)

    password = generate_password(4)
    expired_at = (datetime.datetime.utcnow() + datetime.timedelta(days=7))

    share_id = str(uuid.uuid4())
    add_share(share_id, session_id, latest_interaction_id, thread_json, expired_at, password)
    return jsonify(
        thread=thread,
        expired_at=expired_at.timestamp(),
        password=password,
        share_id=share_id,
        status="new"
    ), 200

@app.route("/sessions/<session_id>/share", methods=['DELETE'])
@cross_origin(origins=["http://localhost:3000"], methods=["POST", "PUT", "OPTIONS", "GET", "DELETE", "PATCH"], all_headers="*", supports_credentials=True, automatic_options=True)
@profile
def clear_share(session_id):
    user_hash = calculate_md5(get_user_identity() + ' ' + request.headers.get('host'))
    if not check_session_owner(session_id, user_hash):
        return jsonify({"status": "error", "message": "Session not found or does not belong to the current user"}), 404

    delete_share(session_id)

    return jsonify({'message': 'Share cleared successfully'})


@app.route("/share/<share_id>", methods=['GET'])
@cross_origin(origins=["http://localhost:3000"], methods=["POST", "PUT", "OPTIONS", "GET", "DELETE", "PATCH"], all_headers="*", supports_credentials=True, automatic_options=True)
def get_share(share_id):
    password = request.args.get('password')
    resp = query_share(share_id, password)
    if resp['status'] == 'error':
        resp['message'] = 'Share not found'
        return jsonify(resp), 400
    return jsonify(json.loads(resp['thread'])), 200

@app.route("/interactions/<interaction_id>/feedback", methods=['POST'])
@cross_origin(origins=["http://localhost:3000"], methods=["POST", "PUT", "OPTIONS", "GET", "DELETE", "PATCH"], all_headers="*", supports_credentials=True, automatic_options=True)
@profile
def feedback_interaction(interaction_id):
    interaction = query_interaction(interaction_id)
    if interaction is None:
        return jsonify({'status': 'error', 'message': 'Interaction Not Found'}), 404

    host = request.headers.get('host')
    user_hash = calculate_md5(get_user_identity() + ' ' + host)
    if not check_session_owner(interaction['session_id'], user_hash):
        return jsonify({"status": "error", "message": "Session not found or does not belong to the current user"}), 404

    interactions = query_interactions(interaction['session_id'])
    messages = get_historical_msgs(interactions, interaction_id) + [{'content': interaction['request'], 'role': 'user'}, {'content': interaction['response'], 'role': 'assistant'}]

    headers = {
        'x-pingchat-access-key': PINGCHAT_ACCESS_KEY,
        'x-pingchat-access-secret': PINGCHAT_ACCESS_SECRET,
        'x-pingchat-user': get_user_identity(),
        'x-trace-id': interaction_id,
    }

    score = request.json.get('score')
    feedback = request.json.get('feedback')

    if score is not None:
        if int(score) not in (-1, 0, 1):
            return jsonify({'status': 'error', 'message': 'Invalid score'}), 400
        update_interaction(interaction_id, {"score": int(score)})
        if int(score) == 1:
            requests.post(f'{CHATBOT_SERVICE_ENDPOINT}/like', json=messages, headers=headers)
        elif int(score) == -1:
            requests.post(f'{CHATBOT_SERVICE_ENDPOINT}/dislike', json={"messages": messages}, headers=headers)
    if feedback is not None:
        if score is not None and int(score) != -1:
            return jsonify({'status': 'error', 'message': 'Invalid feedback'}), 400
        if score is None:
            if interaction is not None:
                if interaction['score'] != -1:
                    return jsonify({'status': 'error', 'message': 'Invalid feedback'}), 400
        delete_feedback(interaction_id)
        feedback_id = str(uuid.uuid4())
        _type = feedback.get('type')
        description = feedback.get('description')
        related_link = feedback.get('related_link')
        insert_feedback(feedback_id, interaction_id, _type, description, related_link)
        requests.post(f'{CHATBOT_SERVICE_ENDPOINT}/dislike', json={'messages': messages, 'type': _type, 'description': description, 'related_link': related_link}, headers=headers)
    return jsonify({'message': 'Feedback recorded'}), 200


@app.route('/linguflow/applications/<application_id>/async_run', methods=['POST'])
@cross_origin(origins=["http://localhost:3000"], methods=["POST", "PUT", "OPTIONS", "GET", "DELETE", "PATCH"], all_headers="*", supports_credentials=True, automatic_options=True)
@profile
def run_linguflow_app(application_id):
    payload = {'input': request.json.get('input')}
    headers = {'x-linguflow-user': get_user_identity(), 'x-linguflow-access-key': LANGLINK_ACCESS_KEY,'x-linguflow-access-secret': LANGLINK_ACCESS_SECRET}
    # The input must be list. It's better to check if the application's input is a list.
    response = requests.post(f'{LINGUFLOW_ENDPOINT}/applications/{application_id}/async_run', json=payload, headers=headers, timeout=600)
    if response.status_code != 200:
        return jsonify({'status': 'error', 'message': response.reason}), 500
    else:
        output = response.json()
        return jsonify(output), 200
    
@app.route('/linguflow/interactions/<interaction_id>', methods=['GET'])
@cross_origin(origins=["http://localhost:3000"], methods=["POST", "PUT", "OPTIONS", "GET", "DELETE", "PATCH"], all_headers="*", supports_credentials=True, automatic_options=True)
@profile
def get_linguflow_interaction(interaction_id):
    headers = {'x-linguflow-user': get_user_identity(), 'x-linguflow-access-key': LANGLINK_ACCESS_KEY,'x-linguflow-access-secret': LANGLINK_ACCESS_SECRET}
    response = requests.get(f'{LINGUFLOW_ENDPOINT}/interactions/{interaction_id}', headers=headers, timeout=600)
    if response.status_code != 200:
        return jsonify({'status': 'error', 'message': response.reason}), 500
    else:
        output = response.json()['interaction']['output']
        return jsonify(output), 200
    
@app.route('/linguflow/interactions/<interaction_id>/scores', methods=['POST'])
@cross_origin(origins=["http://localhost:3000"], methods=["POST", "PUT", "OPTIONS", "GET", "DELETE", "PATCH"], all_headers="*", supports_credentials=True, automatic_options=True)
#@profile
def score_interaction(interaction_id):
    payload = {'value': request.json.get('score'), 'comment': request.json.get('comment')}
    headers = {'x-linguflow-user': get_user_identity(), 'x-linguflow-access-key': LANGLINK_ACCESS_KEY,'x-linguflow-access-secret': LANGLINK_ACCESS_SECRET}
    response = requests.post(f'{LINGUFLOW_ENDPOINT}/interactions/{interaction_id}/scores', json=payload, headers=headers, timeout=600)
    if response.status_code != 200:
        return jsonify({'status': 'error', 'message': response.reason}), 500
    else:
        output = response.json()
        return jsonify(output), 200

@app.route("/ping", methods=['GET'])
@cross_origin(origins=["http://localhost:3000"], methods=["POST", "PUT", "OPTIONS", "GET", "DELETE", "PATCH"], all_headers="*", supports_credentials=True, automatic_options=True)
def ping():
    status, resp = ping_db()
    if status:
        return jsonify({'message': 'db ok'}), 200
    return jsonify({'status': 'error', 'message': resp}), 500

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