from flask import Flask, render_template, request, jsonify
from google import genai
from google.genai import types
import google.generativeai as genai
from PIL import Image
import io
import fitz
import re
import base64
import json
from dotenv import load_dotenv
import os

load_dotenv()  # Loads environment variables from .env file
gemini_api_key = "AIzaSyBtXV2xJbrWVV57B5RWy_meKXOA59HFMeY"
if not gemini_api_key:
    raise ValueError("GEMINI_API_KEY not set in environment variables.")
genai.configure(api_key=gemini_api_key)

app = Flask(__name__)
app.secret_key = "supersecretkey"


def process_file(file):
    """Process file into an appropriate format for Gemini."""
    if file.content_type == 'application/pdf':
        try:
            pdf_bytes = file.read()
            pdf_document = fitz.open(stream=pdf_bytes, filetype="pdf")
            first_page = pdf_document[0]
            pix = first_page.get_pixmap()
            img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
            return img, file.filename
        except Exception as e:
            raise Exception(f"Error processing PDF {file.filename}: {str(e)}")
    else:
        try:
            image_bytes = file.read()
            image = Image.open(io.BytesIO(image_bytes))
            return image, file.filename
        except Exception as e:
            raise Exception(f"Error processing image {file.filename}: {str(e)}")


@app.route("/", methods=["GET", "POST"])
def index():
    if request.method == "POST":
        if "files" not in request.files:
            return jsonify({"error": "No files uploaded"})
        files = request.files.getlist("files")
        if not files or files[0].filename == "":
            return jsonify({"error": "No files selected"})
        
        # Retrieve the selected language from the form
        language = request.form.get("language", "English")
        
        try:
            processed_files = []
            for file in files:
                processed_content, filename = process_file(file)
                processed_files.append((processed_content, filename))
            filenames = [f[1] for f in processed_files]
            files_list = "\n".join([f"- {fname}" for fname in filenames])
            content_parts = []
            for img, _ in processed_files:
                content_parts.append(img)
            # Updated prompt: include the language instruction and refined content details
            content_parts.append(f"""
Please analyze all {len(files)} soil reports and generate a comprehensive soil analysis report in JSON format, entirely in {language}.
The JSON should include exactly five keys: "Overview", "Soil Composition", "Nutrient Levels", "pH Balance & Moisture Content", and "Recommendations & Actionable Steps". 
Each key should have a string value containing the respective analysis.
Reports being analyzed:
{files_list}
""")
            model = genai.GenerativeModel("gemini-1.5-flash")
            response = model.generate_content(content_parts)

            if response and response.text:
                summary = response.text.strip()
                if summary.lower().startswith("please provide"):
                    return jsonify({"error": "Could not analyze the documents"})

                # Optionally, attempt to load a JSON object from the response
                try:
                    analysis_json = json.loads(summary)
                except Exception:
                    # If the response is not a valid JSON, return it as a simple text field.
                    analysis_json = {"analysis": summary}

                return jsonify(analysis_json)
            else:
                return jsonify({"error": "No analysis could be generated"})
        except Exception as e:
            return jsonify({"error": f"Error processing files: {str(e)}"})
    return render_template("index.html")


if __name__ == "__main__":
    port = int(os.environ.get("PORT", 7860))
    app.run(host='0.0.0.0', port=port)