import functools
import datetime
import json
from environs import Env
from sqlalchemy import create_engine, text
from sqlalchemy.exc import OperationalError, TimeoutError
from sqlalchemy.pool import QueuePool
from tenacity import retry, stop_after_attempt, wait_random, retry_if_exception_type

env = Env()
env.read_env()

STORAGE_DSN = env.str("STORAGE_DSN")

db = create_engine(STORAGE_DSN, connect_args={'ssl': {'ssl_ca': '/etc/ssl/certs/ca-certificates.crt'}}, poolclass=QueuePool, pool_size=10, max_overflow=20)


@functools.lru_cache
@retry(stop=stop_after_attempt(3), wait=wait_random(min=1, max=2), retry=retry_if_exception_type((OperationalError, TimeoutError)))
def check_session_owner(session_id, user_hash):
    conn = db.connect()
    query = text("SELECT user_hash FROM sessions WHERE id = :session_id and is_deleted = false")
    result = conn.execute(query, {"session_id": session_id})
    row = result.fetchone()
    conn.close()

    if row is None or row[0] != user_hash:
        return False

    return True


@retry(stop=stop_after_attempt(3), wait=wait_random(min=1, max=2), retry=retry_if_exception_type((OperationalError, TimeoutError)))
def query_sessions(user_hash, offset, limit, order):
    conn = db.connect()
    #query = text("SELECT id, title FROM sessions WHERE user_hash = :user_hash AND is_deleted = false ORDER BY created_at {} LIMIT :limit OFFSET :offset".format(order))
    query = text("""SELECT s.id, s.title, COALESCE(i.latest_updated_at, s.created_at) AS updated_at
FROM sessions s
LEFT JOIN (
  SELECT session_id, MAX(updated_at) AS latest_updated_at
  FROM interactions
  GROUP BY session_id
) i ON s.id = i.session_id
WHERE s.user_hash = :user_hash AND s.is_deleted = false
ORDER BY updated_at {} 
LIMIT :limit OFFSET :offset
    """.format(order))
    result = conn.execute(query, {"user_hash": user_hash, "limit": limit, "offset": offset})
    sessions = [{'id': row[0], 'title': row[1], 'updated_at': row[2]} for row in result.fetchall()]
    conn.close()
    return sessions


@retry(stop=stop_after_attempt(3), wait=wait_random(min=1, max=2), retry=retry_if_exception_type((OperationalError, TimeoutError)))
def count_sessions(user_hash):
    conn = db.connect()
    query = text("SELECT COUNT(*) FROM sessions WHERE user_hash = :user_hash AND is_deleted = false")
    result = conn.execute(query, {"user_hash": user_hash})
    row = result.fetchone()
    conn.close()
    if row is not None:
        return row[0]
    return 0


@retry(stop=stop_after_attempt(3), wait=wait_random(min=1, max=2), retry=retry_if_exception_type((OperationalError, TimeoutError)))
def insert_session(session_id, user_hash, user_email, user_source, created_at, title):
    conn = db.connect()
    query = text("INSERT INTO sessions (id, user_hash, user_email, user_source, created_at, title) VALUES (:session_id, :user_hash, :user_email, :user_source, :created_at, :title)")
    conn.execute(query, {"session_id": session_id, "user_hash": user_hash, "user_email": user_email, "user_source": user_source, "created_at": created_at, "title": title})
    conn.commit()
    conn.close()


@retry(stop=stop_after_attempt(3), wait=wait_random(min=1, max=2), retry=retry_if_exception_type((OperationalError, TimeoutError)))
def query_session(session_id):
    conn = db.connect()
    # query = text("SELECT id, created_at, public_link, expired_at, password, title FROM sessions WHERE id = :session_id AND is_deleted = false")
    query = text("SELECT id, created_at, title FROM sessions WHERE id = :session_id AND is_deleted = false")
    result = conn.execute(query, {"session_id": session_id})
    session_row = result.fetchone()
    conn.close()

    if session_row is None:
        return None

    # session_id, created_at, public_link, expired_at, password, title = session_row
    session_id, created_at, title = session_row

    return {
        "id": session_id,
        "created_at": created_at.timestamp(),
        "title": title
    }


@retry(stop=stop_after_attempt(3), wait=wait_random(min=1, max=2), retry=retry_if_exception_type((OperationalError, TimeoutError)))
def update_session(session_id, payload):
    sql = "UPDATE sessions SET "
    params = {}

    for key, value in payload.items():
        sql += f"{key} = :{key}, "
        params[key] = value

    sql = sql.rstrip(", ")

    sql += " WHERE id = :session_id"
    params['session_id'] = session_id

    conn = db.connect()
    conn.execute(text(sql), params)

    conn.commit()
    conn.close()


@retry(stop=stop_after_attempt(3), wait=wait_random(min=1, max=2), retry=retry_if_exception_type((OperationalError, TimeoutError)))
def is_session_valid_public(session_id):
    conn = db.connect()
    query = text("SELECT id, thread, expired_at, password FROM shares WHERE session_id = :session_id and expired_at > CURRENT_TIMESTAMP()")
    result = conn.execute(query, {"session_id": session_id})
    row = result.fetchone()
    conn.close()
    if row is not None:
        share_id, thread, expired_at, password = row
        return True, {"share_id": share_id, "thread": thread, "expired_at": expired_at.timestamp(), "password": password}
    return False, None


@retry(stop=stop_after_attempt(3), wait=wait_random(min=1, max=2), retry=retry_if_exception_type((OperationalError, TimeoutError)))
def is_interaction_belongs_to_session(session_id, interaction_id):
    conn = db.connect()
    query = text("SELECT id FROM interactions WHERE session_id = :session_id AND id = :interaction_id")
    result = conn.execute(query, {"session_id": session_id, "interaction_id": interaction_id})
    row = result.fetchone()
    conn.close()
    if row is not None:
        return True
    return False


@retry(stop=stop_after_attempt(3), wait=wait_random(min=1, max=2), retry=retry_if_exception_type((OperationalError, TimeoutError)))
def query_interactions(session_id):
    conn = db.connect()
    query = text("SELECT id, created_at, updated_at, parent_id, request, response, score, error, links FROM interactions WHERE session_id = :session_id ORDER BY updated_at DESC")
    result = conn.execute(query, {"session_id": session_id})
    interaction_rows = result.fetchall()
    conn.close()

    interactions = {}
    for interaction_row in interaction_rows:
        interaction_id, created_at, updated_at, parent_id, request, response, score, error, links = interaction_row

        interactions[interaction_id] = {
            "id": interaction_id,
            "parent_id": parent_id,
            "request": request,
            "response": response,
            "score": score,
            "error": error,
            "updated_at": updated_at.timestamp(),
            "created_at": created_at.timestamp(),
            "links": json.loads(links)
        }

    return interactions


@retry(stop=stop_after_attempt(3), wait=wait_random(min=1, max=2), retry=retry_if_exception_type((OperationalError, TimeoutError)))
def insert_interaction(id, session_id, created_at, updated_at, parent_id, request, response, revised_q, recalls, links, score, error):
    conn = db.connect()
    query = text("INSERT INTO interactions (id, session_id, created_at, updated_at, parent_id, request, response, revised_q, recalls, links, score, error) VALUES (:id, :session_id, :created_at, :updated_at, :parent_id, :request, :response, :revised_q, :recalls, :links, :score, :error)")
    conn.execute(
        query, 
        {
            "id": id, "session_id": session_id, "created_at": created_at, "updated_at": updated_at,
            "parent_id": parent_id, "request": request, "response": response, "revised_q": revised_q,
            "recalls": json.dumps(recalls), "links": json.dumps(links), "score": score, "error": error
        }
    )
    conn.commit()
    conn.close()


@retry(stop=stop_after_attempt(3), wait=wait_random(min=1, max=2), retry=retry_if_exception_type((OperationalError, TimeoutError)))
def update_interaction(interaction_id, payload):
    sql = "UPDATE interactions SET "
    params = {}

    for key, value in payload.items():
        sql += f"{key} = :{key}, "
        params[key] = value

    sql = sql.rstrip(", ")

    sql += " WHERE id = :interaction_id"
    params['interaction_id'] = interaction_id

    conn = db.connect()
    conn.execute(text(sql), params)
    
    conn.commit()
    conn.close()


@retry(stop=stop_after_attempt(3), wait=wait_random(min=1, max=2), retry=retry_if_exception_type((OperationalError, TimeoutError)))
def query_interaction(interaction_id):
    conn = db.connect()
    query = text("SELECT id, session_id, created_at, updated_at, parent_id, request, response, revised_q, recalls, links, score, error FROM interactions WHERE id = :interaction_id")
    result = conn.execute(query, {"interaction_id": interaction_id})
    interaction_row = result.fetchone()
    conn.close()
    if interaction_row is None:
        return None
    interaction_id, session_id, created_at, updated_at, parent_id, request, response, revised_q, recalls, links, score, error = interaction_row
    return {
        "id": interaction_id,
        "session_id": session_id,
        "created_at": created_at.timestamp(),
        "updated_at": updated_at.timestamp(),
        "parent_id": parent_id,
        "request": request,
        "response": response,
        "revised_q": revised_q,
        "recalls": json.loads(recalls),
        "links": json.loads(links),
        "score": score,
        "error": error
    }


@retry(stop=stop_after_attempt(3), wait=wait_random(min=1, max=2), retry=retry_if_exception_type((OperationalError, TimeoutError)))
def find_latest_interaction(session_id):
    conn = db.connect()
    query = text("SELECT id, updated_at FROM interactions WHERE session_id = :session_id ORDER BY updated_at DESC LIMIT 1")
    result = conn.execute(query, {"session_id": session_id})
    latest_interaction = result.fetchone()
    conn.close()
    if latest_interaction is None:
        return None, None
    return latest_interaction[0], latest_interaction[1].timestamp()


@retry(stop=stop_after_attempt(3), wait=wait_random(min=1, max=2), retry=retry_if_exception_type((OperationalError, TimeoutError)))
def add_share(share_id, session_id, latest_interaction_id, thread_json, expired_at, password):
    conn = db.connect()
    query = text("INSERT INTO shares (id, session_id, last_interaction_id, thread, expired_at, password) VALUES(:id, :session_id, :last_interaction_id, :thread, :expired_at, :password) ON DUPLICATE KEY UPDATE id=:id, session_id=:session_id, last_interaction_id=:last_interaction_id, thread=:thread, expired_at=:expired_at, password=:password")
    conn.execute(query, {"id": share_id, "session_id": session_id, "last_interaction_id": latest_interaction_id, "thread": thread_json, "expired_at": expired_at, "password": password})
    conn.commit()
    conn.close()


@retry(stop=stop_after_attempt(3), wait=wait_random(min=1, max=2), retry=retry_if_exception_type((OperationalError, TimeoutError)))
def delete_share(session_id):
    conn = db.connect()
    query = text("DELETE FROM shares WHERE session_id = :session_id")
    result = conn.execute(query, {"session_id": session_id})
    affected_rows = result.rowcount
    conn.commit()
    conn.close()
    if affected_rows == 0:
        return None
    return None


@retry(stop=stop_after_attempt(3), wait=wait_random(min=1, max=2), retry=retry_if_exception_type((OperationalError, TimeoutError)))
def query_share(share_id, password_input):
    conn = db.connect()
    query = text("SELECT thread, expired_at, password FROM shares WHERE id = :id")
    res = conn.execute(query, {"id": share_id})
    row = res.fetchone()
    conn.close()
    if row is None:
        return {"status": "error", "message": "not found"}
    thread, expired_at, password = row
    if expired_at < datetime.datetime.utcnow():
        return {"status": "error", "message": "expired"}
    if password_input != password:
        return {"status": "error", "message": "wrong password"}
    return {"status": "ok", "thread": thread}


@retry(stop=stop_after_attempt(3), wait=wait_random(min=1, max=2), retry=retry_if_exception_type((OperationalError, TimeoutError)))
def insert_feedback(feedback_id, interaction_id, _type, description, correct_link):
    created_at = datetime.datetime.utcnow()
    conn = db.connect()
    query = text("INSERT INTO feedbacks (id, interaction_id, type, description, correct_link, created_at) VALUES (:id, :interaction_id, :type, :description, :correct_link, :created_at)")
    conn.execute(query, {"id": feedback_id, "interaction_id": interaction_id, "type": _type, "description": description, "correct_link": correct_link, "created_at": created_at})
    conn.commit()
    conn.close()


@retry(stop=stop_after_attempt(3), wait=wait_random(min=1, max=2), retry=retry_if_exception_type((OperationalError, TimeoutError)))
def delete_feedback(interaction_id):
    conn = db.connect()
    query = text("DELETE FROM feedbacks WHERE interaction_id = :interaction_id")
    result = conn.execute(query, {"interaction_id": interaction_id})
    affected_rows = result.rowcount
    conn.commit()
    conn.close()
    if affected_rows == 0:
        return None
    return None


@retry(stop=stop_after_attempt(3), wait=wait_random(min=1, max=2), retry=retry_if_exception_type((OperationalError, TimeoutError)))
def ping_db():
    try:
        conn = db.connect()
        conn.connection.ping()
        conn.close()
        return True, None
    except OperationalError as e:
        return False, str(e)