import os
import json
from flask import Flask, render_template, request, jsonify, redirect, url_for
from werkzeug.utils import secure_filename
from huggingface_hub import InferenceClient
import pandas as pd
import docx
from PyPDF2 import PdfReader

app = Flask(__name__)

# Set up file upload configurations
UPLOAD_FOLDER = "uploads"
app.config["UPLOAD_FOLDER"] = UPLOAD_FOLDER
ALLOWED_EXTENSIONS = {"txt", "csv", "json", "pdf", "docx"}

# Retrieve Hugging Face API key securely from environment variables
api_key = os.getenv("HF_API_KEY")
if not api_key:
    raise ValueError("Hugging Face API key not found. Set 'HF_API_KEY' in your Space secrets.")

# Initialize Hugging Face Inference Client
client = InferenceClient(api_key=api_key)


# Function to check allowed file types
def allowed_file(filename):
    return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS


# Function to read uploaded files and extract content
def extract_file_content(filepath, file_type):
    content = ""
    try:
        if file_type == "txt":
            with open(filepath, "r", encoding="utf-8") as file:
                content = file.read()
        elif file_type == "csv":
            df = pd.read_csv(filepath)
            content = df.to_string()
        elif file_type == "json":
            with open(filepath, "r", encoding="utf-8") as file:
                content = json.dumps(json.load(file), indent=4)
        elif file_type == "pdf":
            reader = PdfReader(filepath)
            content = "".join(page.extract_text() for page in reader.pages)
        elif file_type == "docx":
            doc = docx.Document(filepath)
            content = "\n".join(paragraph.text for paragraph in doc.paragraphs)
    except Exception as e:
        raise ValueError(f"Error extracting file content: {e}")
    return content


# Function to send content to Hugging Face model
def get_bot_response(prompt):
    try:
        response = client.text_generation(
            prompt=prompt,
            model="Qwen/Qwen2.5-Coder-32B-Instruct",
            max_tokens=500
        )
        return response
    except Exception as e:
        return f"Error in model response: {e}"


# Route: Home Page (File Upload Form)
@app.route("/", methods=["GET", "POST"])
def upload_file():
    if request.method == "POST":
        # Check if file is uploaded
        if "file" not in request.files:
            return jsonify({"error": "No file part"}), 400

        file = request.files["file"]

        if file.filename == "":
            return jsonify({"error": "No selected file"}), 400

        if file and allowed_file(file.filename):
            filename = secure_filename(file.filename)
            filepath = os.path.join(app.config["UPLOAD_FOLDER"], filename)
            os.makedirs(app.config["UPLOAD_FOLDER"], exist_ok=True)
            file.save(filepath)

            # Extract file content
            file_type = filename.rsplit(".", 1)[1].lower()
            try:
                content = extract_file_content(filepath, file_type)
            except Exception as e:
                return jsonify({"error": str(e)}), 500

            # Send content to Hugging Face model
            response = get_bot_response(content)

            return jsonify({"response": response})

        else:
            return jsonify({"error": "File type not allowed"}), 400

    return render_template("upload.html")


# Route: Retrieve Model Response (API Endpoint)
@app.route("/generate", methods=["POST"])
def generate_response():
    data = request.get_json()
    prompt = data.get("prompt")

    if not prompt:
        return jsonify({"error": "No prompt provided"}), 400

    # Send prompt to Hugging Face model
    response = get_bot_response(prompt)
    return jsonify({"response": response})


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