from flask import Flask, render_template, request, jsonify
import os
import requests
import json
import http.server
import socketserver
import threading
import subprocess
import time
from dotenv import load_dotenv
from PIL import Image
import io
from transformers import pipeline
import base64

app = Flask(__name__)
load_dotenv()

# Configuration
INSTAGRAM_ACCESS_TOKEN = os.getenv('INSTAGRAM_ACCESS_TOKEN')
INSTAGRAM_USER_ID = os.getenv('INSTAGRAM_USER_ID')
HUGGINGFACE_API_TOKEN = os.getenv('HUGGINGFACE_API_TOKEN')
UPLOAD_FOLDER = 'uploads'
PORT = 8000

# Ensure upload folder exists
os.makedirs(UPLOAD_FOLDER, exist_ok=True)

# Start local HTTP server to serve images
Handler = http.server.SimpleHTTPRequestHandler
httpd = socketserver.TCPServer(("", PORT), Handler)
server_thread = threading.Thread(target=httpd.serve_forever)
server_thread.daemon = True
server_thread.start()

# Start ngrok to create a public URL
ngrok_process = None

def start_ngrok():
    global ngrok_process
    ngrok_process = subprocess.Popen(["ngrok", "http", str(PORT)], stdout=subprocess.PIPE)
    time.sleep(2)  # Wait for ngrok to start
    resp = requests.get("http://localhost:4040/api/tunnels")
    public_url = resp.json()["tunnels"][0]["public_url"]
    return public_url

# Hugging Face captioning pipeline
caption_pipeline = pipeline("image-to-text", model="Salesforce/blip-image-captioning-base")

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

@app.route('/upload', methods=['POST'])
def upload_image():
    try:
        # Get base64 image from request
        image_data = request.json['image'].split(',')[1]
        image_bytes = base64.b64decode(image_data)
        
        # Save image to uploads folder
        image_path = os.path.join(UPLOAD_FOLDER, 'captured_image.jpg')
        with open(image_path, 'wb') as f:
            f.write(image_bytes)
        
        # Generate caption using Hugging Face
        image = Image.open(io.BytesIO(image_bytes))
        caption_result = caption_pipeline(image)
        caption = caption_result[0]['generated_text']

        # Get public URL from ngrok
        public_url = start_ngrok()
        image_url = f"{public_url}/uploads/captured_image.jpg"

        # Create Instagram media container
        create_url = f"https://graph.facebook.com/v20.0/{INSTAGRAM_USER_ID}/media"
        payload = {
            "image_url": image_url,
            "caption": caption,
            "access_token": INSTAGRAM_ACCESS_TOKEN
        }
        response = requests.post(create_url, params=payload)
        result = response.json()

        if 'id' in result:
            creation_id = result['id']
            # Publish the media
            publish_url = f"https://graph.facebook.com/v20.0/{INSTAGRAM_USER_ID}/media_publish"
            publish_payload = {
                "creation_id": creation_id,
                "access_token": INSTAGRAM_ACCESS_TOKEN
            }
            publish_response = requests.post(publish_url, params=publish_payload)
            publish_result = publish_response.json()

            if 'id' in publish_result:
                return jsonify({"status": "success", "message": "Image posted to Instagram", "caption": caption})
            else:
                return jsonify({"status": "error", "message": publish_result.get('error', {}).get('message', 'Failed to publish')})
        else:
            return jsonify({"status": "error", "message": result.get('error', {}).get('message', 'Failed to create media container')})
    except Exception as e:
        return jsonify({"status": "error", "message": str(e)})
    finally:
        # Clean up ngrok
        if ngrok_process:
            ngrok_process.terminate()

@app.route('/shutdown', methods=['POST'])
def shutdown():
    httpd.shutdown()
    if ngrok_process:
        ngrok_process.terminate()
    return jsonify({"status": "success", "message": "Server shutdown"})

if __name__ == '__main__':
    try:
        app.run(debug=True)
    finally:
        httpd.shutdown()
        if ngrok_process:
            ngrok_process.terminate()