import time
import re
from flask import Flask, request, jsonify
from flask_cors import CORS
from google.oauth2 import id_token
from google.auth.transport import requests as google_requests
import logging
import psycopg2
from psycopg2.extras import RealDictCursor
from model import process_query_with_mistral

_real_time = time.time
time.time = lambda: _real_time() - 1

# Database connection parameters
DATABASE_CONFIG = {
    "dbname": "HealthAIDB",
    "user": "postgres",
    "password": "Oleg2005",
    "host": "postgres",
    "port": 5432,
}

import logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

try:
    conn = psycopg2.connect(**DATABASE_CONFIG)
    logger.info("Database connection established successfully")
except Exception as e:
    logger.error(f"Error connecting to database: {e}", exc_info=True)
    conn = None

def init_db():
    create_users_query = """
    CREATE TABLE IF NOT EXISTS users (
    id SERIAL PRIMARY KEY,
    name TEXT NOT NULL,
    email TEXT UNIQUE NOT NULL,
    google_id TEXT,
    password TEXT,
    phone TEXT,
    role TEXT,
    bio TEXT,
    picture TEXT
    );
    """
    create_chat_history_query = """
    CREATE TABLE IF NOT EXISTS chat_history (
        id SERIAL PRIMARY KEY,
        user_email TEXT NOT NULL,
        chat TEXT NOT NULL,
        user_data TEXT,
        created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
    );
    """
    try:
        with conn.cursor() as cur:
            cur.execute(create_users_query)
            cur.execute(create_chat_history_query)
            conn.commit()
        logger.info("Database tables initialized successfully")
    except Exception as e:
        logger.error(f"Error initializing database tables: {e}", exc_info=True)
        conn.rollback()

if conn:
    init_db()

app = Flask(__name__)
CORS(app, resources={r"/api/*": {"origins": "*"}})

CLIENT_ID = "532143017111-4eqtlp0oejqaovj6rf5l1ergvhrp4vao.apps.googleusercontent.com"

def save_user_to_db(name, email, google_id=None, password=None):
    logger.info(f"Saving user {name} with email: {email} to the database")
    try:
        with conn.cursor(cursor_factory=RealDictCursor) as cur:
            cur.execute(
                """
                INSERT INTO users (name, email, google_id, password)
                VALUES (%s, %s, %s, %s)
                ON CONFLICT (email) DO NOTHING
                """,
                (name, email, google_id, password)
            )
            conn.commit()
        logger.info(f"User {name} ({email}) saved successfully")
    except Exception as e:
        logger.error(f"Error saving user {name} ({email}) to database: {e}", exc_info=True)

@app.route('/api/verify', methods=['POST'])
def verify_token():
    logger.info("Received token verification request")
    data = request.get_json()
    token = data.get('token')
    if not token:
        logger.warning("Token not provided in request")
        return jsonify({'error': 'No token provided'}), 400
    try:
        id_info = id_token.verify_oauth2_token(token, google_requests.Request(), CLIENT_ID)
        user_email = id_info.get('email')
        user_name = id_info.get('name')
        google_id = id_info.get('sub')
        logger.info(f"Token verified for user: {user_name} ({user_email})")
        save_user_to_db(name=user_name, email=user_email, google_id=google_id)
        return jsonify({'message': 'Authentication successful', 'user': {'email': user_email, 'name': user_name}}), 200
    except ValueError as e:
        logger.error(f"Token verification error: {e}", exc_info=True)
        return jsonify({'error': 'Invalid token'}), 400


@app.errorhandler(404)
def not_found(e):
    return jsonify({'error': 'Not found'}), 404

@app.errorhandler(500)
def internal_error(e):
    return jsonify({'error': 'Internal server error'}), 500


@app.route('/api/register', methods=['POST'])
def register():
    logger.info("Received new user registration request")
    data = request.get_json()
    name = data.get('name')
    email = data.get('email')
    password = data.get('password')
    if not all([name, email, password]):
        logger.warning("Not all required fields provided for registration")
        return jsonify({'error': 'All fields are required'}), 400
    try:
        with conn.cursor(cursor_factory=RealDictCursor) as cur:
            cur.execute("SELECT * FROM users WHERE email = %s", (email,))
            existing_user = cur.fetchone()
            if existing_user:
                logger.warning(f"User with email {email} already exists")
                return jsonify({'error': 'User already exists'}), 409
        save_user_to_db(name=name, email=email, password=password)
        logger.info(f"User {name} ({email}) registered successfully")
        return jsonify({'message': 'User registered successfully'}), 201
    except Exception as e:
        logger.error(f"Error during user registration: {e}", exc_info=True)
        return jsonify({'error': str(e)}), 500

@app.route('/api/login', methods=['POST'])
def login():
    logger.info("Received login request")
    data = request.get_json()
    email = data.get('email')
    password = data.get('password')
    if not all([email, password]):
        logger.warning("Email or password not provided")
        return jsonify({'error': 'Email and password are required'}), 400
    try:
        with conn.cursor(cursor_factory=RealDictCursor) as cur:
            cur.execute("SELECT * FROM users WHERE email = %s", (email,))
            user = cur.fetchone()
            if not user or user.get('password') != password:
                logger.warning(f"Invalid credentials for email: {email}")
                return jsonify({'error': 'Invalid credentials'}), 401
        logger.info(f"User {user.get('name')} ({email}) logged in successfully")
        return jsonify({'message': 'Login successful', 'user': {'name': user.get('name'), 'email': user.get('email')}}), 200
    except Exception as e:
        logger.error(f"Error during user login: {e}", exc_info=True)
        return jsonify({'error': str(e)}), 500


@app.route('/api/update_profile', methods=['PUT'])
def update_profile():
    data = request.get_json()
    email = data.get('email')
    if not email:
        return jsonify({'error': 'Email is required'}), 400

    # Fields to update; if not provided, the current value remains.
    name = data.get('name')
    phone = data.get('phone')
    role = data.get('role')
    bio = data.get('bio')
    picture = data.get('picture')

    try:
        with conn.cursor(cursor_factory=RealDictCursor) as cur:
            cur.execute(
                """
                UPDATE users
                SET name = COALESCE(%s, name),
                    phone = COALESCE(%s, phone),
                    role = COALESCE(%s, role),
                    bio = COALESCE(%s, bio),
                    picture = COALESCE(%s, picture)
                WHERE email = %s
                """,
                (name, phone, role, bio, picture, email)
            )
            conn.commit()
        logger.info(f"Profile updated for {email}")
        return jsonify({'message': 'Profile updated successfully'}), 200
    except Exception as e:
        logger.error(f"Error updating profile: {e}")
        return jsonify({'error': str(e)}), 500

@app.route('/api/chat', methods=['POST'])
def chat():
    logger.info("Received chat request")
    data = request.get_json()
    query = data.get('query', '')
    user_email = data.get('email')
    chat_id = data.get('chatId')

    if not query:
        logger.warning("No query provided")
        return jsonify({'error': 'No query provided'}), 400

    logger.info(f"Processing request for chatId: {chat_id if chat_id else 'new chat'} | Query: {query}")

    # Retrieve chat context from the database
    chat_context = ""
    if chat_id:
        try:
            with conn.cursor(cursor_factory=RealDictCursor) as cur:
                cur.execute("SELECT chat, user_data FROM chat_history WHERE id = %s", (chat_id,))
                result = cur.fetchone()
                if result:
                    chat_context = result.get("chat", "")
                    logger.info(f"Loaded chat context for chatId {chat_id}: {chat_context}")
                else:
                    logger.info(f"No chat context found for chatId {chat_id}")
        except Exception as e:
            logger.error(f"Error loading chat context from DB: {e}", exc_info=True)

    logger.info("Calling process_query_with_mistral function")
    response_obj = process_query_with_mistral(query, chat_id=chat_id, chat_context=chat_context)
    best_answer = response_obj.get("best_answer", "") if isinstance(response_obj, dict) else str(response_obj)
    logger.info(f"Response from process_query_with_mistral: {best_answer}")

    best_answer = re.sub(r'[*#]', '', best_answer)
    best_answer = re.sub(r'(\d\.\s)', r'\n\n\1', best_answer)
    best_answer = re.sub(r':\s-', r':\n-', best_answer)

    # Update or create chat_history record including user_data if available
    if chat_id:
        try:
            with conn.cursor(cursor_factory=RealDictCursor) as cur:
                cur.execute("SELECT chat FROM chat_history WHERE id = %s", (chat_id,))
                existing_chat = cur.fetchone()
                if existing_chat:
                    updated_chat = existing_chat['chat'] + f"\nUser: {query}\nBot: {best_answer}"
                    if "patient_data" in response_obj:
                        cur.execute("UPDATE chat_history SET chat = %s, user_data = %s WHERE id = %s",
                                    (updated_chat, response_obj["patient_data"], chat_id))
                    else:
                        cur.execute("UPDATE chat_history SET chat = %s WHERE id = %s", (updated_chat, chat_id))
                    conn.commit()
                    logger.info(f"Chat history for chatId {chat_id} updated successfully")
                else:
                    with conn.cursor(cursor_factory=RealDictCursor) as cur2:
                        cur2.execute(
                            "INSERT INTO chat_history (user_email, chat) VALUES (%s, %s) RETURNING id",
                            (user_email, f"User: {query}\nBot: {best_answer}")
                        )
                        new_chat_id = cur2.fetchone()['id']
                        conn.commit()
                        chat_id = new_chat_id
                        logger.info(f"New chat created with chatId: {chat_id}")
        except Exception as e:
            logger.error(f"Error updating/creating chat history: {e}", exc_info=True)
            return jsonify({'error': str(e)}), 500
    else:
        try:
            with conn.cursor(cursor_factory=RealDictCursor) as cur:
                cur.execute(
                    "INSERT INTO chat_history (user_email, chat) VALUES (%s, %s) RETURNING id",
                    (user_email, f"User: {query}\nBot: {best_answer}")
                )
                new_chat_id = cur.fetchone()['id']
                conn.commit()
                chat_id = new_chat_id
                logger.info(f"New chat created with chatId: {chat_id}")
        except Exception as e:
            logger.error(f"Error creating new chat: {e}", exc_info=True)
            return jsonify({'error': str(e)}), 500

    return jsonify({'response': {'best_answer': best_answer, 'model': response_obj.get("model", ""), 'chatId': chat_id}}), 200

@app.route('/api/save_user_data', methods=['POST'])
def save_user_data():
    logger.info("Received request to save user data")
    data = request.get_json()
    chat_id = data.get('chatId')
    user_data = data.get('userData')
    if not chat_id or not user_data:
        return jsonify({'error': 'chatId and userData are required'}), 400
    try:
        with conn.cursor(cursor_factory=RealDictCursor) as cur:
            cur.execute("UPDATE chat_history SET user_data = %s WHERE id = %s", (user_data, chat_id))
            conn.commit()
        logger.info(f"User data for chatId {chat_id} updated successfully")
        return jsonify({'message': 'User data updated successfully'}), 200
    except Exception as e:
        logger.error(f"Error updating user data: {e}", exc_info=True)
        return jsonify({'error': str(e)}), 500

@app.route('/api/chat_history', methods=['GET'])
def get_chat_history():
    logger.info("Received request to get chat history")
    user_email = request.args.get('email')
    if not user_email:
        return jsonify({'error': 'User email is required'}), 400
    try:
        with conn.cursor(cursor_factory=RealDictCursor) as cur:
            cur.execute(
                "SELECT id, chat, user_data, created_at FROM chat_history WHERE user_email = %s ORDER BY created_at DESC",
                (user_email,)
            )
            history = cur.fetchall()
        return jsonify({'history': history}), 200
    except Exception as e:
        logger.error(f"Error getting chat history for {user_email}: {e}", exc_info=True)
        return jsonify({'error': str(e)}), 500

@app.route('/api/chat_history', methods=['DELETE'])
def delete_chat():
    logger.info("Received request to delete chat")
    chat_id = request.args.get('id')
    if not chat_id:
        return jsonify({'error': 'Chat id is required'}), 400
    try:
        with conn.cursor() as cur:
            cur.execute("DELETE FROM chat_history WHERE id = %s", (chat_id,))
            conn.commit()
        return jsonify({'message': 'Chat deleted successfully'}), 200
    except Exception as e:
        logger.error(f"Error deleting chat with chatId {chat_id}: {e}", exc_info=True)
        return jsonify({'error': str(e)}), 500

@app.route('/api/get_user_data', methods=['GET'])
def get_user_data():
    chat_id = request.args.get('chatId')
    if not chat_id:
        return jsonify({'error': 'Chat id is required'}), 400
    try:
        with conn.cursor(cursor_factory=RealDictCursor) as cur:
            cur.execute("SELECT user_data FROM chat_history WHERE id = %s", (chat_id,))
            result = cur.fetchone()
        if result and result.get("user_data"):
            return jsonify({'user_data': result.get("user_data")}), 200
        else:
            return jsonify({'user_data': None}), 200
    except Exception as e:
        logger.error(f"Error retrieving user data: {e}", exc_info=True)
        return jsonify({'error': str(e)}), 500

@app.route('/api/chat_history_detail', methods=['GET'])
def chat_history_detail():
    logger.info("Received request to get chat details")
    chat_id = request.args.get('id')
    if not chat_id:
        return jsonify({'error': 'Chat id is required'}), 400
    try:
        with conn.cursor(cursor_factory=RealDictCursor) as cur:
            cur.execute("SELECT id, chat, user_data, created_at FROM chat_history WHERE id = %s", (chat_id,))
            chat = cur.fetchone()
        if not chat:
            return jsonify({'error': 'Chat not found'}), 404
        return jsonify({'chat': chat}), 200
    except Exception as e:
        logger.error(f"Error getting chat details for chatId {chat_id}: {e}", exc_info=True)
        return jsonify({'error': str(e)}), 500

if __name__ == '__main__':
    logger.info("Starting Flask application")
    app.run(host='0.0.0.0', port=5000, debug=True)