import os
import uuid
import time
import json
import base64
import openai
import pinecone
import functools
import pandas as pd
from barely_json import parse
from typing import List, Dict
from neo4j import GraphDatabase
from authing import AuthenticationClient
from sqlalchemy import create_engine, text
from flask import Flask, request, session, redirect


app = Flask(__name__)

TEXT_EMBEDDINGS_MODEL = 'text-embedding-ada-002'
PINECONE_API_KEY = os.environ["PINECONE_API_KEY"]
NEO4J_ENDPOINT = os.environ["NEO4J_ENDPOINT"]
NEO4J_PASSWORD = os.environ["NEO4J_PASSWORD"]
OSSINSIGHT_DSN = os.environ["OSSINSIGHT_DSN"]
STORAGE_DSN = os.environ["STORAGE_DSN"]
openai.api_key = os.environ["OPENAI_API_KEY"]
app.secret_key = os.environ["APP_SECRET_KEY"]
AUTH_APP_ID = os.environ["AUTH_APP_ID"]
AUTH_APP_HOST = os.environ["AUTH_APP_HOST"]
AUTH_APP_SECRET = os.environ["AUTH_APP_SECRET"]
AUTH_REDIRECT_URI = os.environ["AUTH_REDIRECT_URI"]

pc = pinecone.Pinecone(api_key=PINECONE_API_KEY)
pinecone_index = pc.Index(host="https://ossinsight-5c794ec.svc.aped-4627-b74a.pinecone.io")
auth_client = client = AuthenticationClient(
    AUTH_APP_ID,
    AUTH_APP_HOST,
    AUTH_APP_SECRET,
    redirect_uri=AUTH_REDIRECT_URI,
)
neo4j_driver = GraphDatabase.driver(NEO4J_ENDPOINT, auth=('neo4j', NEO4J_PASSWORD))
store = create_engine(STORAGE_DSN, connect_args={'ssl': {'ssl_ca': '/etc/ssl/certs/ca-certificates.crt'}})
engine = create_engine(OSSINSIGHT_DSN, max_overflow=-1, connect_args={'ssl': {'ssl_ca': '/etc/ssl/certs/ca-certificates.crt'}})

def profile(handler):
    @functools.wraps(handler)
    def ensure_profile(*args, **kwargs):
        if 'profile' in session:
            return handler(*args, **kwargs)
        return {
            'auth_url': auth_client.build_authorize_url(scope="profile openid", state=request.headers.get('state'))
        }, 401
    return ensure_profile

def get_embedding(text: str, model: str) -> List[float]:
    result = openai.Embedding.create(
        model=model,
        input=text,
    )
    return result["data"][0]["embedding"]

# the function used to get embedding
def get_text_embedding(text: str) -> List[float]:
    return get_embedding(text, TEXT_EMBEDDINGS_MODEL)

def sync(ids: List[int]) -> None:
    with neo4j_driver.session(database="neo4j") as session:
        df = session.run(
            "match (n:Example) where id(n) in $ids return id(n) as id, properties(n) as metadata",
            ids=ids,
        ).to_df()
        if len(df) > 0:
            pinecone_index.upsert(
                vectors=[
                    {
                        'id': str(uuid.UUID(fields=(0, 0, 0, 0, 0, r.get('id')))),
                        'values': get_text_embedding(r.get('metadata')['Q']),
                        'metadata': r.get('metadata'),
                    } for _, r in df.iterrows()
                ],
                namespace='example'
            )

        df = session.run(
            "match (n:Entity) where id(n) in $ids return id(n) as id, properties(n) as metadata",
            ids=ids,
        ).to_df()
        if len(df) > 0:
            pinecone_index.upsert(
                vectors=[
                    {
                        'id': str(uuid.UUID(fields=(0, 0, 0, 0, 0, r.get('id')))),
                        'values': get_text_embedding(r.get('metadata')['name']),
                        'metadata': r.get('metadata'),
                    } for _, r in df.iterrows()
                ],
                namespace='entity'
            )
        
        df = session.run(
            'match (n:Column) where id(n) in $ids return id(n) as id, properties(n) as metadata',
            ids=ids,
        ).to_df()
        if len(df) > 0:
            pinecone_index.upsert(
                vectors=[
                    {
                        'id': str(uuid.UUID(fields=(0, 0, 0, 0, 0, r.get('id')))),
                        'values': get_text_embedding(r.get('metadata')['desc']),
                        'metadata': r.get('metadata'),
                    } for _, r in df.iterrows()
                ],
                namespace='column'
            )

def set_entity(entity: str, desc: str, source: str = None) -> None:
    set_source = ', e.source = $source'
    if source is None:
        set_source = ''
        
    with neo4j_driver.session(database="neo4j") as session:
        rs = session.run(
            "merge (e:Entity { name: $entity }) "
            f"on match set e.ts = $ts, e.desc = $desc{set_source} "
            f"on create set e.ts = $ts, e.desc = $desc{set_source} "
            "return id(e) as id",
            entity=entity, desc=desc, ts=int(time.time()), source=source,
        )
        sync([r.get('id') for r in rs])

def set_example(question: str, answer: str, source: str = None) -> None:
    set_source = ', e.source = $source'
    if source is None:
        set_source = ''
        
    with neo4j_driver.session(database="neo4j") as session:
        rs = session.run(
            "merge (e:Example { Q: $question }) "
            f"on match set e.ts = $ts, e.A = $answer{set_source} "
            f"on create set e.ts = $ts, e.A = $answer{set_source} "
            "return id(e) as id",
            question=question, answer=answer, ts=int(time.time()), source=source,
        )
        sync([r.get('id') for r in rs])

def set_column(table: str, column: str, desc: str, source: str = None) -> None:
    set_source = ', c.source = $source'
    if source is None:
        set_source = ''

    with neo4j_driver.session(database="neo4j") as session:
        rs = session.run(
            "merge (c:Column { table: $table, column: $column }) "
            f"on match set c.ts = $ts, c.desc = $desc{set_source} "
            f"on create set c.ts = $ts, c.desc = $desc{set_source} "
            "return id(c) as id",
            table=table, column=column, ts=int(time.time()), desc=desc, source=source,
        )
        sync([r.get('id') for r in rs])
        
def link_column(entity: str, table: str, column: str, source: str = None) -> None:
    set_source = ', r.source = $source'
    if source is None:
        set_source = ''
        
    with neo4j_driver.session(database="neo4j") as session:
        session.run(
            "match (e:Entity { name: $entity }) "
            "match (c:Column { table: $table, column: $column }) "
            "merge (e)-[r:REFER]->(c) "
            f"on match set r.ts = $ts{set_source} "
            f"on create set r.ts = $ts{set_source} ",
            entity=entity, table=table, column=column, ts=int(time.time()), source=source,
        )
            
def unique_top_k(xs: List[Dict], top_k: int) -> List[Dict]:
    d = {}
    for x in xs:
        if x['id'] in d and x['score'] < d[x['id']]['score']:
            continue
        d[x['id']] = x
    return sorted(d.values(), key=lambda x: x['score'], reverse=True)[:top_k]

def recall(questions: List[str], top_k: int = 5) -> Dict[str, List]:
    assert type(questions) == list
    
    query_embeddings = [get_text_embedding(question) for question in questions]
    
    entities = unique_top_k(sum([pinecone_index.query(
        namespace='entity', 
        top_k=top_k,
        include_metadata=True, 
        vector=query_embedding,
    )['matches'] for query_embedding in query_embeddings], []), top_k)

    examples = unique_top_k(sum([pinecone_index.query(
        namespace='example', 
        top_k=5,
        include_metadata=True, 
        vector=query_embedding,
    )['matches'] for query_embedding in query_embeddings], []), top_k)
    
    ids = [uuid.UUID(x['id']).node for x in entities]
    with neo4j_driver.session(database="neo4j") as session:
        rs = session.run(
            "match (e:Entity)-->(c:Column) where id(e) in $ids "
            "return distinct c.table as table, c.column as column, c.desc as desc", 
            ids=ids
        )
        columns = [{
            'table': r.get('table'),
            'column': r.get('column'),
            'desc': r.get('desc'),
        } for r in rs]
        
    return {
        'entities': [x['metadata'] for x in entities], 
        'examples': [x['metadata'] for x in examples], 
        'columns': columns,
    }

def prompt(question, r):    
    column_str = '\n\n'.join([c['table'] + '.' + c['column'] + ': ' + c['desc'] for c in r['columns']])
    entity_str = '\n\n'.join([c['desc'] for c in r['entities']])
    example_str = '\n\n'.join(['Q: ' + e['Q'] + '\nA: ' + e['A'] for e in r['examples']])
    
    return '''# MySQL SQL
# Table Schema
Table github_events, columns = [id, type, created_at, repo_id, repo_name, actor_id, actor_login, additions, deletions, action, number, org_login, org_id, state, closed_at, comments, pr_merged_at, pr_merged, pr_changed_files, pr_review_comments, pr_or_issue_id, push_size, push_distinct_size, creator_user_login, creator_user_id, pr_or_issue_created_at]
- Column type, enums = ['PullRequestEvent', 'PushEvent', 'IssueCommentEvent', 'IssuesEvent', 'PullRequestReviewCommentEvent', 'WatchEvent', 'CreateEvent', 'DeleteEvent', 'ForkEvent', 'ReleaseEvent']
- Column action:
* type in [PullRequestReviewCommentEvent, IssueCommentEvent, ReviewEvent]: created
* type in [PullRequestEvent, IssuesEvent]: opened, closed, reopened
- Column number, number is issue number
- Column created_at, closed_at, pr_merged_at, pr_or_issue_created_at DEFAULT '1970-01-01 00:00:00'
- Column additions, deletions: only existed in PullRequestEvent
- Column actor_id, actor_login: who trigger the event

Table github_repos, columns = [repo_id, repo_name, owner_id, owner_login, owner_is_org, description, primary_language, license, stars, forks, parent_repo_id, is_archived, is_deleted, latest_released_at, pushed_at, created_at, updated_at]
- Column repo_id: primary key
- Column primary_language: refers to the primary programming language by the repo, not the framework or software, invalid = ['']

Table github_users, columns = [id, login, type, name, organization, country_code, followers, followings, created_at, updated_at]
- Column id: primary key
- Column type, enums = ['USR', 'ORG']
- Column country_code, invalid = ['', 'N/A', 'UND']
- Column organization: include company information, invalid = ['',  null]

Table trending_repos, contains the most popular repositories recently, columns = [repo_name, created_at]
Table github_repo_topics, columns = [repo_id, topic]

# Table Relations
github_events.repo_id = github_repos.repo_id
github_events.actor_id = github_users.id
github_events.actor_login = github_users.login
github_events.creator_user_id = github_users.id
github_repos.owner_id = github_users.id
github_repos.repo_id = github_repo_topics.repo_id
trending_repos.repo_name = github_repos.repo_name

# Common Sense
''' + entity_str + '''

# Example
''' + example_str + '''

# ChartOptions
type Column = string; // must be the column name in the SQL result!!!
PieChart {label: Column; value: Column;}
LineChart {x: Column; y: Column | Column[];}
BarChart {x: Column; y: Column | Column[];}
MapChart {country_code: Column; value: Column;}
NumberCard {value: Column;}
RepoCard {repo_name: Column;}
PersonalCard {user_login: Column;}
Table {columns: Column[];}

# ChartConfig
type ChartName = 'PieChart' | 'LineChart' | 'BarChart' | 'MapChart' | 'NumberCard' | 'RepoCard' | 'PersonalCard' | 'Table';
interface ChartConfig {chartName: ChartName; title: string; options: ChartOptions;}

use the answer.json template:
{
    "RQ": "The question after revised, marked as RQ, RQ should be a clear and concise statement of the question",
    "keywords": ["Output every key information('who', 'when', 'what', 'time range') in RQ"],
    "links": ["List the GitHub links of the repos mentioned in RQ"],
    "subQuestions": ["List at most 2 important sub-questions with No in order (e.g: No1; No2; ...) in array to answer the RQ"],
    "notClear": "What is not clear or needs to narrow down",
    "assumption": "What's your assumptions about RQ and sub-questions",
    "deps": "Dependency graph of each sub-question, Question Number only",
    "CQ": "Integrate (RQ,assumptions) into a clarified question,then diff_words_from(RQ) and mark with (<b>)",
    "sqlCanAnswer": boolean,
    "sql": "use the best practice to write a single SQL to answer the the question",
    "chart": ChartConfig
}

Begin:
Let's think step by step, generate an answer.json to answer the question "In GitHub, ''' + question + '''" (Do not overthink):
---
answer.json // must only output a pretty JSON!!!
---'''

def generate(p):
    completion = openai.ChatCompletion.create(
        model="gpt-3.5-turbo",
        messages=[{
            "role": "user", 
            "content": p,
        }],
        temperature=0.0,
        top_p=1.0,
        max_tokens=1024,
        frequency_penalty=0.0,
        presence_penalty=0.0,
    )
    return completion.choices[0].message.content

def get_sql(answer):
    x = parse(answer)
    if x.get('sql'):
        return x['sql']
    elif x.get('Sql'):
        return x['Sql']
    elif x.get('SQL'):
        return x['SQL']
    else:
        raise ValueError('No valid parsed SQL')
    
def execute(question_id):
    with engine.connect() as conn:
        question = conn.execute(
            text('select title from explorer_questions where TO_BASE64(id) = :id'),
            id=question_id
        ).fetchone()[0]
    
    r = recall([question])
    for k, v in r.items():
        yield({k: json.dumps(v)})
        
    p = prompt(question, r)
    yield {'prompt': p}
    
    answer = generate(p)
    yield {'answer': answer}    
    yield {'revised_title': parse(answer).get('RQ')}
    
    sql = get_sql(answer)
    yield {'query_sql': sql}
    
    df = pd.read_sql_query(text(sql), engine)
    yield {'result': json.dumps({
        'fields': [{'name': c} for c in df.columns],
        'rows': json.loads(df[:100].to_json(orient='records')),
    })}

@app.route("/ping")
def ping():
    return {'status': 'success'}
    
@app.route("/execute/manual", methods=['POST'])
def execute_manual():
    question_id = request.json['id']
    r = {
        'id': question_id,
        'status': None,
        'error': None,
        'result': None,
        'prompt': None,
        'answer': None,
        'revised_title': None,
        'query_sql': None,
        'entities': None,
        'examples': None,
        'columns': None,
    }
    
    try:
        for x in execute(question_id):
            for k, v in x.items():
                r[k] = v
        r['status'] = 'success'
    except Exception as e:
        r['status'] = 'error'
        r['error'] = str(e)
        
    with store.connect() as conn:
        statement = text(f"""
            REPLACE INTO `manual-run`({", ".join(r.keys())})
            VALUES({", ".join([f':{c}' for c in r.keys()])})
        """)

        conn.execute(statement, r)

    return r

@app.route("/execute/daily", methods=['POST'])
def execute_daily():
    question_id = request.json['id']
    r = {
        'id': question_id, 
        'status': None,
        'error': None,
        'result': None,
        'prompt': None,
        'answer': None,
        'revised_title': None,
        'query_sql': None,
        'entities': None,
        'examples': None,
        'columns': None,
    }
    
    try:
        for x in execute(question_id):
            for k, v in x.items():
                r[k] = v
        r['status'] = 'success'
    except Exception as e:
        r['status'] = 'error'
        r['error'] = str(e)
        
    with store.connect() as conn:
        statement = text(f"""
            REPLACE INTO `daily-run`({", ".join(r.keys())})
            VALUES({", ".join([f':{c}' for c in r.keys()])})
        """)

        conn.execute(statement, r)

    return r

@app.route("/questions/<path:path>", methods=['GET'])
def questions(path: str):
    if '-' in path:
        df = pd.read_sql_query(
            text("select * from explorer_questions where date(created_at) = :day"), 
            engine, params={ 'day': path }
        )
        if len(df) == 0:
            return []
    else:
        df = pd.read_sql_query(
            text("select * from explorer_questions where TO_BASE64(id) = :id"), 
            engine, params={ 'id': path }
        )
        if len(df) == 0:
            return {}
    
    df['id'] = df.id.apply(lambda x: base64.b64encode(x).decode('utf-8'))

    statement = 'select * from `daily-run` where id in ({})'.format(', '.join([f'"{x}"' for x in df.id]))
    daily_df = pd.read_sql_query(statement, store).add_prefix('daily_')
    if not daily_df.empty:
        df = pd.merge(df, daily_df, left_on='id', right_on='daily_id', how='left')

    statement = 'select * from `manual-run` where id in ({})'.format(', '.join([f'"{x}"' for x in df.id]))
    manual_df = pd.read_sql_query(statement, store).add_prefix('manual_')
    if not manual_df.empty:
        df = pd.merge(df, manual_df, left_on='id', right_on='manual_id', how='left')

    if '-' in path:
        return json.loads(df.to_json(orient='records'))
    else:
        return json.loads(df.to_json(orient='records'))[0]

@app.route("/trend", methods=["GET"])
def trend():
    df = pd.read_sql_query(
        "select date(created_at) as time, count(*) as total from explorer_questions group by date(created_at)", 
        engine,
    )
    err_df = pd.read_sql_query(
        '''
        select date(created_at) as time, count(*) as error from explorer_questions 
        where status = 'error' group by date(created_at)
        ''',
        engine,
    )
    return json.loads(pd.merge(df, err_df, left_on='time', right_on='time').to_json(orient='records'))

@app.route("/prompts", methods=['POST'])
def prompts():
    question = request.json['question']
    r = recall([question])
    return r

@app.route("/login", methods=['GET'])
def login():
    code = request.args.get('code')
    if not code:
        return {'status': 'error', 'message': 'no code in query params'}, 400
    r = auth_client.get_access_token_by_code(code)
    if 'access_token' not in r:
        return {'status': 'error', 'message': 'invalid access token'}, 400
    user_info = auth_client.get_user_info_by_access_token(r['access_token'])
    session['profile'] = user_info['profile']
    state = request.args.get('state')
    if not state or not state.startswith('http'):
        return {'status': 'success'}
    return redirect(state)

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

@app.route("/examples", methods=['PUT'])
@profile
def modify_example():
    Q = request.json['Q']
    A = request.json['A']
    
    set_example(Q, A, 'user')
    
    return {
        'status': 'success',
        'Q': Q,
        'A': A,
    }

@app.route("/entities", methods=['PUT'])
@profile
def modify_entity():
    entity = request.json['name']
    desc = request.json['desc']
    
    set_entity(entity, desc, 'user')
    
    return {
        'status': 'success',
        'name': entity,
        'desc': desc,
    }