from flask import Flask, jsonify, request
from webscout import WEBS
from webscout.LLM import LLM

app = Flask(__name__)

@app.route('/search', methods=['GET'])
def WEBScout2_search_search():
    query = request.args.get('query')
    if not query:
        return jsonify({'error': 'Query parameter missing'})

    limit = request.args.get('limit', default=10, type=int) # Default limit to 10 if not provided
    WEBS_instance = WEBS() # Instantiate WEBS without context manager
    responses = []
    for i, r in enumerate(WEBS_instance.text(query, region='wt-wt', safesearch='off', timelimit='y')):
        if i == limit: # Use the custom limit
            break
        responses.append(r)
    return jsonify(responses)


@app.route('/mws', methods=['POST'])
def WEBScout_search():
    data = request.get_json()
    if 'query' not in data:
        return jsonify({'error': 'Query parameter missing'})

    query = data['query']
    WEBS_instance = WEBS() # Instantiate WEBS without context manager
    responses = []
    for i, r in enumerate(WEBS_instance.text(query, region='wt-wt', safesearch='off', timelimit='y')):
        if i == 2: # Limiting the results to 2
            break
        responses.append(r)
    return jsonify(responses)

@app.route('/health', methods=['GET'])
def health_check():
    return jsonify({'status': 'ok'})

@app.route('/video', methods=['GET'])
def WEBScout_videos():
    params = request.args
    if 'keywords' not in params:
        return jsonify({'error': 'Keywords parameter missing'})

    keywords = params['keywords']
    WEBS_instance = WEBS() # Instantiate WEBS without context manager
    responses = []
    for r in WEBS_instance.videos(
        keywords,
        region="wt-wt",
        safesearch="off",
        timelimit="w",
        resolution="high",
        duration="medium",
        max_results=10,
    ):
        responses.append(r)
    return jsonify(responses)

@app.route('/img', methods=['GET'])
def WEBScout2_images():
    params = request.args
    if 'keywords' not in params:
        return jsonify({'error': 'Keywords parameter missing'})

    keywords = params['keywords']
    WEBS_instance = WEBS() # Instantiate WEBS without context manager
    responses = []
    for r in WEBS_instance.images(
        keywords,
        region="wt-wt",
        safesearch="off",
        size=None,
        type_image=None,
        layout=None,
        license_image=None,
        max_results=10,
    ):
        responses.append(r)
    return jsonify(responses)

@app.route('/news', methods=['GET'])
def WEBScout_news():
    params = request.args
    if 'keywords' not in params:
        return jsonify({'error': 'Keywords parameter missing'})
    keywords = params['keywords']
    WEBS_instance = WEBS() # Instantiate WEBS without context manager
    responses = []
    for r in WEBS_instance.news(
        keywords,
        region="wt-wt",
        safesearch="off",
        timelimit="m",
        max_results=10
    ):
        responses.append(r)
    return jsonify(responses)


@app.route('/translate', methods=['GET'])
def WEBScout_translate():
    params = request.args
    if 'keywords' not in params or 'to' not in params:
        return jsonify({'error': 'Keywords or target language parameter missing'})
    keywords = params['keywords']
    target_language = params['to']
    WEBS_instance = WEBS() # Instantiate WEBS without context manager
    translation = WEBS_instance.translate(keywords, to=target_language)
    return jsonify(translation)

@app.route('/chat', methods=['POST'])
def chat_gpt():
    user_input = request.json.get('message')
    if user_input is None:
        return jsonify({'error': 'Message parameter missing'})

    llm = LLM("mistralai/Mistral-7B-Instruct-v0.1") # Initialize the LLM class with a model
    response = llm.mistral_chat([{"role": "user", "content": user_input}]) # Call the mistral_chat method
    return jsonify({"response": response})

if __name__ == '__main__':
    app.run(debug=True)