import os
import zipfile
import shutil
import json
from flask import Flask, render_template, request, send_from_directory, session, redirect, url_for
import google.generativeai as genai
from werkzeug.utils import secure_filename
from flask_session import Session
from dotenv import load_dotenv

# Load .env variables
load_dotenv()
GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")

# Config Gemini
genai.configure(api_key=GEMINI_API_KEY)
model = genai.GenerativeModel('gemini-1.5-pro-latest')

# Folders
UPLOAD_FOLDER = 'uploads'
EXTRACT_FOLDER = 'uploads_extracted'
OUTPUT_FOLDER = 'outputs'
ALLOWED_EXTENSIONS = {'txt', 'pdf', 'docx', 'py', 'js', 'html', 'css', 'json', 'java', 'cpp', 'c', 'md', 'zip'}

# Flask app
app = Flask(__name__)
app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
app.config['EXTRACT_FOLDER'] = EXTRACT_FOLDER
app.config['OUTPUT_FOLDER'] = OUTPUT_FOLDER
app.secret_key = 'supersecretkey'
app.config['SESSION_TYPE'] = 'filesystem'
Session(app)

# Ensure folders exist
os.makedirs(UPLOAD_FOLDER, exist_ok=True)
os.makedirs(EXTRACT_FOLDER, exist_ok=True)
os.makedirs(OUTPUT_FOLDER, exist_ok=True)

# Helpers
def allowed_file(filename):
    return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS

def read_file(filepath):
    try:
        with open(filepath, 'r', encoding='utf-8', errors='ignore') as f:
            return f.read()
    except Exception as e:
        return f"ERROR READING FILE: {str(e)}"

def save_txt(content, filename):
    filepath = os.path.join(EXTRACT_FOLDER, filename)
    os.makedirs(os.path.dirname(filepath), exist_ok=True)
    with open(filepath, "w", encoding="utf-8") as f:
        f.write(content)

def extract_zip(zip_path, extract_to):
    with zipfile.ZipFile(zip_path, 'r') as zip_ref:
        zip_ref.extractall(extract_to)

def get_project_tree():
    tree = []
    for root, dirs, files in os.walk(EXTRACT_FOLDER):
        for name in files:
            rel_path = os.path.relpath(os.path.join(root, name), EXTRACT_FOLDER)
            tree.append(rel_path)
    return tree

# Routes
@app.route('/', methods=['GET', 'POST'])
def index():
    if request.method == 'POST':
        # Clear previous project
        shutil.rmtree(EXTRACT_FOLDER)
        os.makedirs(EXTRACT_FOLDER)

        files = request.files.getlist('files')
        for file in files:
            if file and allowed_file(file.filename):
                filename = secure_filename(file.filename)
                filepath = os.path.join(UPLOAD_FOLDER, filename)
                file.save(filepath)

                if filename.endswith('.zip'):
                    extract_zip(filepath, EXTRACT_FOLDER)
                else:
                    save_txt(read_file(filepath), filename)

        session['chat_history'] = []
        session['ai_understood'] = True
        session['project_tree'] = get_project_tree()

        return redirect(url_for('chat'))

    return render_template('index.html')

@app.route('/chat', methods=['GET', 'POST'])
def chat():
    if 'ai_understood' not in session or not session['ai_understood']:
        return redirect(url_for('index'))

    ai_reply = ''
    if request.method == 'POST':
        user_message = request.form.get('user_message')
        chat_history = session.get('chat_history', [])

        # Build project context
        all_text = ''
        for file_path in get_project_tree():
            full_path = os.path.join(EXTRACT_FOLDER, file_path)
            content = read_file(full_path)
            all_text += f"\n--- File: {file_path} ---\n{content}\n"

        # Build chat prompt
        chat_prompt = ''
        for chat in chat_history:
            chat_prompt += f"User: {chat['user']}\nAI: {chat['ai']}\n"

        full_prompt = f"""
You are an expert AI project code editor.
The user uploaded this project:

{all_text}

Chat history:
{chat_prompt}

Now user says:
{user_message}

Please reply with:
- what files you edited (if any)
- content of new/edited files
- if user says DELETE file, remove it
- reply in readable format.

Example format:

### Edited Files:
File: filename.py
Content:
<content here>

### New Files:
File: newfile.py
Content:
<content here>

### Deleted Files:
filename1
filename2

### Notes:
<summary notes here>
"""

        # Run Gemini API
        response = model.generate_content(full_prompt)
        ai_reply = response.text

        # Parse AI reply to save files
        if "### New Files:" in ai_reply:
            parts = ai_reply.split("### New Files:")[-1].split("###")[0].strip().split("File:")
            for part in parts[1:]:
                lines = part.strip().split("\n")
                filename = lines[0].strip()
                content = "\n".join(lines[2:]).strip()
                save_txt(content, filename)

        if "### Edited Files:" in ai_reply:
            parts = ai_reply.split("### Edited Files:")[-1].split("###")[0].strip().split("File:")
            for part in parts[1:]:
                lines = part.strip().split("\n")
                filename = lines[0].strip()
                content = "\n".join(lines[2:]).strip()
                save_txt(content, filename)

        if "### Deleted Files:" in ai_reply:
            parts = ai_reply.split("### Deleted Files:")[-1].split("###")[0].strip().split("\n")
            for filename in parts:
                filename = filename.strip()
                if filename:
                    path_to_delete = os.path.join(EXTRACT_FOLDER, filename)
                    if os.path.exists(path_to_delete):
                        os.remove(path_to_delete)

        # Save chat history
        chat_history.append({'user': user_message, 'ai': ai_reply})
        session['chat_history'] = chat_history
        session['project_tree'] = get_project_tree()

    return render_template('chat.html',
                           chat_history=session.get('chat_history', []),
                           project_tree=session.get('project_tree', []),
                           ai_reply=ai_reply)

@app.route('/download_project')
def download_project():
    # Zip project
    zip_filename = 'project.zip'
    zip_path = os.path.join(OUTPUT_FOLDER, zip_filename)
    shutil.make_archive(os.path.splitext(zip_path)[0], 'zip', EXTRACT_FOLDER)
    return send_from_directory(OUTPUT_FOLDER, zip_filename, as_attachment=True)

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