from flask import Flask, request, render_template, redirect, url_for, session, send_from_directory
from flask_talisman import Talisman
from functools import wraps
import os
from security import validate_user
from data_handler import save_link, save_file, retrieve_uploads, handle_download, get_file_path

app = Flask(__name__, template_folder='../templates')
app.secret_key = os.urandom(24)  # Generate a more secure secret key
talisman = Talisman(app, content_security_policy={
    'default-src': ["'self'"],
    'script-src': ["'self'", "'unsafe-inline'"]  # Allow inline scripts
})

UPLOAD_DIRECTORY = "../assets"
if not os.path.exists(UPLOAD_DIRECTORY):
    os.makedirs(UPLOAD_DIRECTORY)

# Login required decorator
def login_required(f):
    @wraps(f)
    def decorated_function(*args, **kwargs):
        if 'username' not in session:
            print("User is not logged in, redirecting to login page")
            return redirect(url_for('login', next=request.url))
        print("User is logged in, proceeding to requested page")
        return f(*args, **kwargs)
    return decorated_function

@app.before_request
def ensure_login():
    if 'username' not in session and request.endpoint not in ('login', 'static'):
        return redirect(url_for('login'))

@app.route('/')
@login_required
def index():
    return render_template("index.html")

@app.route('/login', methods=['GET', 'POST'])
def login():
    if request.method == 'POST':
        username = request.form['username']
        password = request.form['password']
        valid, message = validate_user(username, password)
        
        if valid:
            session['username'] = username
            return redirect(url_for('index'))
        else:
            return render_template("login.html", error=message)

    return render_template("login.html")

@app.route('/logout')
def logout():
    session.clear()
    return redirect(url_for('login'))

@app.route('/upload/link', methods=['POST'])
@login_required
def upload_link():
    link = request.form['link']
    uploader = session['username']
    save_link(uploader, link)
    return redirect(url_for('index'))

@app.route('/upload/file', methods=['POST'])
@login_required
def upload_file():
    if 'file' not in request.files:
        return redirect(url_for('index'))
    
    file = request.files['file']
    uploader = session['username']
    save_file(uploader, file)
    return redirect(url_for('index'))

@app.route('/uploads')
@login_required
def view_uploads():
    uploads = retrieve_uploads()
    
    # Categorizing uploads
    links = [upload for upload in uploads if upload[2] == 'link']
    videos = [upload for upload in uploads if upload[2] == 'file' and upload[3].lower().endswith(('.mp4', '.mkv', '.avi'))]
    photos = [upload for upload in uploads if upload[2] == 'file' and upload[3].lower().endswith(('.jpg', '.jpeg', '.png', '.gif'))]
    misc = [upload for upload in uploads if upload[2] == 'file' and upload not in videos + photos]

    return render_template(
        "uploads.html", 
        links=links, 
        videos=videos, 
        photos=photos, 
        misc=misc, 
        username=session['username']
    )

@app.route('/download_link/<int:link_id>', methods=['GET'])
@login_required
def download_link(link_id):
    uploader = session['username']
    upload = handle_download(link_id, uploader)

    if upload[2] == 'link':
        link_content = upload[3]

        # Create a unique filename
        x = 1
        while os.path.exists(f"link_{x}.txt"):
            x += 1
        filename = f"link_{x}.txt"

        # Save the link content to the file
        with open(filename, 'w') as f:
            f.write(link_content)

        # Serve the file
        return send_from_directory(directory=os.getcwd(), filename=filename, as_attachment=True)

@app.route('/download_all_links', methods=['GET'])
@login_required
def download_all_links():
    links = [upload for upload in retrieve_uploads() if upload[2] == 'link']

    if len(links) > 1:
        with open("links_data.txt", 'w') as f:
            for link in links:
                f.write(link[3] + "\n")

        return send_from_directory(directory=os.getcwd(), filename="links_data.txt", as_attachment=True)
    else:
        return redirect(url_for('view_uploads'))

@app.route('/download/<int:upload_id>', methods=['GET'])
@login_required
def download(upload_id):
    uploader = session['username']
    upload = handle_download(upload_id, uploader)
    
    if upload[2] == 'link':
        return f"<a href='{upload[3]}' target='_blank'>{upload[3]}</a>"
    elif upload[2] == 'file':
        # Integrate get_file_path here
        file_path = get_file_path(upload[3])
        return send_from_directory(os.path.dirname(file_path), os.path.basename(file_path), as_attachment=True)

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5000, ssl_context='adhoc')