# Import the necessary functions from ollama, Flask, requests, threading
from ollama import Client
from flask import Flask, request, jsonify, send_from_directory, render_template, session, make_response
from flask_cors import CORS, cross_origin # CORS stands for Cross-Origin Resource Sharing. This is necessary to allow the frontend to make requests to our backend.
import requests
import json
import logging
import os
import utils
from utils import GlobalState

# Create a logger for this module
global_state = GlobalState() # Import the singleton that holds global states (e.g., logger)
logger = global_state.getLogger(__name__)  # Logger for this module, inherit properties of the root logger


# Find out the path to current directory according to the Python interpreter (venv)
logger.debug("Current working directory: %s", os.getcwd())

# Initialize a Flask application
app = Flask(__name__)
app.config['STATIC_FOLDER'] = 'static' # Adjust if needed 

# Set the secret key for session management
secret_key = os.urandom(24)
app.config['SECRET_KEY'] = secret_key # When do I need this. How is it retained between sessions?

# Optionally set other configuration options
app.config['SESSION_PERMANENT'] = False  # Session will expire after each request
app.config['SESSION_TYPE'] = 'filesystem'  # Store sessions on the filesystem


logger.debug("flask app template folder: %s", app.template_folder) 

@app.route('/')
def index():
    """
    This route serves index.html to connecting clients
    """

    session['chat_history'] = []   # The session object (actually, a dictonary) holds the chat session
    logger.debug("Entering route '/'")
    # api_endpoint = os.environ['BE_API_ENDPOINT']  # Retrieve the environment variable 
    api_endpoint = global_state.get_backend_api_ep()  # Retrieve the environment variable 
    logger.debug("Backend API endpoint: %s", api_endpoint)
    use_model = global_state.get_llm()
    with open('smartassist/src/html/client.html', 'r') as f:
        client_html = f.read()
    logger.debug("Client HTML (first few characters): %s", client_html[:50]) # Print to see if it's loading
    # logger.debug("Client HTML (all characters): %s", client_html) # Print to see if it's loading

    return render_template('index.html', api_endpoint=api_endpoint, use_model = use_model, client_content=client_html)

@app.route('/set_session')
def set_session():
    resp = make_response()
    resp.set_cookie('session', 'some-value', samesite='None', secure=True) # Add SameSite attribute here
    return resp

# @app.route('/profile')
# def profile():
#     # Retrieve data from the session
#     user_id = session.get('user_id')

#     if user_id:
#         return f'User ID: {user_id}'
#     else:
#         return 'No user ID found'
    

@app.route('/<path:filename>')
def serve_static(filename):
    return send_from_directory(app.config['STATIC_FOLDER'], filename)


# CORS(app, resources={
#     r"/api/chat": {
#         "origins": "*",
#         "headers": ["Origin", "Content-Type", "Authorization"],
#     }
# })

CORS(app, resources={
    r"/api/chat": {
        "origins": "*"    
    }
})
    
@app.route('/api/tags', methods=['GET'])
def tag(url = "http://localhost:11434/api/tags", headers = None):
# def tag(url = "http://localhost:11434/api/tags", headers = {"Content-Type": "application/json"}):
    """Get a list of models for the server located at url."""
    try:
        logger.debug(f"url: {url} headers: {headers}")
        response = requests.get(url, headers=headers)
        return response.json()
        # return response
    except requests.exceptions.RequestException as e:
        logger.error("Request Exception: %s", str(e))        
        return {'error': 'Failed to process request'}


@app.route('/api/chat', methods=['POST'])
def chat(model = "phi3:mini"):
# def chat(url_server = "http://localhost:11434/api/generate", model = "phi3:mini"):
    """
    This function handles the chat. The frontend client (web browser) calls the 
    backend server through this endpoint (/api/chat) that manage queries
    to the LLM (Large Language Model) server and it also manages the response 
    from the LLM server.
    """
    # Get the message from the JSON in the request body
    data = request.get_json()
    message = data.get('query')
    url_server = data.get('url_server', "https://ollama-test.wara-ops.org/api/generate")  # Use provided URL or default
    model = data.get('model', model)  # Use provided model or default if not provided
    
    # Get chat history from session storage (e.g., a dictionary)
    chat_history = session.get('chat_history', [])
    
    # Add the new message to the chat history
    chat_history.append({'role': 'user', 'message': message})
    
    # Update the session with the new chat history
    session['chat_history'] = chat_history
    
    # Create the data dictionary with chat history
    data_to_send = {
        "model": model,
        'prompt': '\n'.join([f"{item['role']}: {item['message']}" for item in chat_history]),
        "stream": False           
    }
    
    try:
        url = url_server
        headers = {
                    "Content-Type": "application/json",
                    "Authorization": "Basic ZWNzanBlcjoxM2JjMTU4ZDhmNmY5YTU4YTkzZDNmY2I="
        }
        logger.debug(f"url: {url} headers: {headers}")
        response = requests.post(url,
                                 headers=headers, 
                                 data=json.dumps(data_to_send))
        response.raise_for_status()  # Raise an exception for bad status codes
        llm_response = response.json()['response']   # Assuming the LLM's response is under 'response' key
        chat_history.append({'role': 'assistant', 'message': llm_response})  # Add assistant response to chat history
        logger.debug(f"Chat History: {chat_history}")
        return response.json()
    except requests.exceptions.RequestException as e:
        logger.error("Request Exception: %s", str(e))        
        return jsonify({'error': 'Failed to process request'}), 500
    except json.JSONDecodeError as e:
        logger.error("JSON Decode Error: %s", str(e))   # Corresponds to print(f"JSON Decode Error: {e}")             
        return jsonify({'error': 'Invalid JSON response from server'}), 500



@app.route('/smartassist', methods=["POST"])
def smartassist():
    # Extract the query from the incoming JSON data
    data = request.json
    user_query = data['query']

    # Get the response from the OLLAMA API based on the user's query
    # NOTE: Should we append message history here? Maybe interact with SQLlite?
    response = get_response(user_query)

    # Return the response as a JSON object in the HTTP response
    return jsonify({"response": response})

def get_response(user_query):
    client = Client()   # Create a client object for interacting with OLLAMA API
    response = client.generate_response(user_query)   # Generate and retrieve the response based on user's query
    return response     


def run_flask(fport=5005):
    """
    Starts the Flask server
    """
    # Flask endpoint for user interaction
    logger.debug("Entering run_flask()")
    # app.run(port = str(str(fport)), debug=False)
    app.run(port = str(str(fport)), debug=True)
    # app.run(port=5000, debug=True, use_reloader=False)
    logger.debug("Exiting run_flask()")


if __name__ == '__main__':
    # Run the Flask application
    run_flask()