import os
import uuid
import time
import json
import yaml
import torch
import openai
import asyncio
import logging
import tiktoken
import functools
import pinecone
import pandas as pd
from environs import Env
from logs import log_handler
from datetime import datetime
from observability import trace
from response import ResponseChain
from namespace import get_ns_by_tags
from pingchat import load_pingchat_chain
from embedding import PingChatEmbeddings
from callback import QueueCallbackHandler
from urllib.parse import urlparse, urlunparse
from langchain.chat_models import ChatOpenAI
from flask import Flask, request, session, redirect, g

from auth import (
    BaseAuthenticationClient, 
    PingCAPAuthenticationClient, 
    CommunityAuthenticationClient,
)


torch.set_num_threads(os.cpu_count())

env = Env()
env.read_env()

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

PINECONE_API_KEY = env.str('PINECONE_API_KEY')
PINECONE_ENVIRONMENT = env.str('PINECONE_ENVIRONMENT')
PINECONE_INDEX_NAME = env.str('PINECONE_INDEX_NAME', default='embedding-768')
PINECONE_DOCUMENT_VERSION = env.str('PINECONE_DOCUMENT_VERSION', default='v2023-04-24')
AUTH_APP_ID = env.str('AUTH_APP_ID')
AUTH_APP_HOST = env.str("AUTH_APP_HOST")
AUTH_APP_SECRET = env.str("AUTH_APP_SECRET")
AUTH_DEFAULT_SCOPE = env.str("AUTH_DEFAULT_SCOPE", default="profile openid")
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")
OPENAI_API_KEY = env.str("OPENAI_API_KEY")
ENABLE_AUTH = env.bool("ENABLE_AUTH", default=True)

pinecone.init(api_key=PINECONE_API_KEY, environment=PINECONE_ENVIRONMENT)
embedding = PingChatEmbeddings('/data/embedding')
llm = ChatOpenAI(streaming=True, temperature=0, openai_api_key=OPENAI_API_KEY)

@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,
    )

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')

@functools.lru_cache
def get_source_by_access_key(access_key: str) -> str:
    with open('/data/sources.yaml') as f:
        sources = yaml.load(f, Loader=yaml.CLoader)
    return sources.get(access_key, 'pingchat')
    
def get_user_source():
    access_key = request.headers.get('x-pingchat-access-key')
    return get_source_by_access_key(access_key)
    
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

@trace(name="chatbot")
def generate(messages, doc_ns, callback):
    chain = load_pingchat_chain(env, llm, embedding, doc_ns)
    outputs = chain(inputs={'messages': messages})
    r = {
        'content': outputs['content'],
        'links': outputs['links'],
    }
    callback({
        'request': messages,
        'response': {
            **r,
            'recalls': outputs['recalls'],
            'instruction': outputs['instruction'],
        },
    })
    return r

async def generate_stream(messages, doc_ns, callback):
    queue = asyncio.Queue()
    chain = load_pingchat_chain(env, llm, embedding, doc_ns)

    task = asyncio.create_task(trace(
        name="chatbot-stream",
        input_fn = lambda args, kwargs: kwargs["inputs"],
        output_fn = lambda outputs: {
            'content': outputs['content'],
            'links': outputs['links'],
        },
    )(chain.acall)(
        inputs={'messages': messages},
        callbacks=[QueueCallbackHandler(queue, ResponseChain.__name__)],
    ))

    while True:
        t = await queue.get()
        if t is None or isinstance(t, Exception):
            break
        yield json.dumps({
            'content': t,
        }) + '\n'

    outputs = await task

    yield json.dumps({'links': outputs['links']}) + '\n'

    callback({
        'request': messages,
        'response': {
            'content': outputs['content'],
            'links': outputs['links'],
            'recalls': outputs['recalls'],
            'instruction': outputs['instruction'],
        },
    })

@app.before_request
def add_trace_id():
    g.trace_id = request.headers.get('x-trace-id', str(uuid.uuid4()))

@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("/me", methods=['GET'])
@profile
def me():
    return {
        'profile': get_user_identity(),
    }

@app.route("/chat", methods=['POST'])
@profile
def chat():
    ns_tags = request.args.to_dict(flat=False)
    return generate(
        messages=request.json,
        doc_ns=get_ns_by_tags(ns_tags),
        callback=lambda x: app.logger.info({
            **x, 
            'user': get_user_identity(),
            'source': get_user_source(),
            'tags': ns_tags,
        }),
    )

@app.route("/chat-stream", methods=['POST'])
@profile
def chat_stream():
    ns_tags = request.args.to_dict(flat=False)
    user = get_user_identity()
    source = get_user_source()
    trace_id = g.trace_id

    ag = generate_stream(
        request.json,
        get_ns_by_tags(ns_tags),
        lambda x: app.logger.info({
            **x, 
            'user': user,
            'source': source,
            'tags': ns_tags,
            'x-trace-id': trace_id,
        }),
    )
            
    def gen():
        try:
            loop = asyncio.new_event_loop()
            while True:
                yield loop.run_until_complete(
                    asyncio.wait_for(ag.__anext__(), 60),
                )
        except StopAsyncIteration:
            pass
        finally:
            loop.close()
    return app.response_class(gen(), mimetype='text/plain')

@app.route("/like", methods=['POST'])
@profile
def like():
    app.logger.info({
        'messages': request.json,
        'feedback': 'like',
        'user': get_user_identity(),
    })
    return {'status': 'success'}

@app.route("/dislike", methods=['POST'])
@profile
def dislike():
    app.logger.info({
        'messages': request.json,
        'feedback': 'dislike',
        'user': get_user_identity(),
    })
    return {'status': 'success'}

@app.route("/ping")
def ping():
    return {'status': 'success'}