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
import socket
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
httpd = None
ngrok_process = None

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

# Check if port is available
def is_port_available(port):
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
        return s.connect_ex(('localhost', port)) != 0

# Start local HTTP server to serve images
def start_http_server():
    global httpd
    Handler = http.server.SimpleHTTPRequestHandler
    port = PORT
    max_attempts = 5
    for i in range(max_attempts):
        if is_port_available(port):
            try:
                httpd = socketserver.TCPServer(("", port), Handler)
                server_thread = threading.Thread(target=httpd.serve_forever)
                server_thread.daemon = True
                server_thread.start()
                print(f"HTTP server started on port {port}")
                return port
            except Exception as e:
                print(f"Failed to start HTTP server on port {port}: {e}")
                raise
        else:
            print(f"Port {port} is in use, trying port {port + 1}")
            port += 1
    raise OSError("No available ports found after multiple attempts")

# Start ngrok to create a public URL
def start_ngrok(port):
    global ngrok_process
    try:
        ngrok_process = subprocess.Popen(["ngrok", "http", str(port)], stdout=subprocess.PIPE, stderr=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"]
        print(f"Ngrok tunnel started at {public_url}")
        return public_url
    except Exception as e:
        print(f"Failed to start ngrok: {e}")
        raise

# Hugging Face captioning pipeline
try:
    caption_pipeline = pipeline("image-to-text", model="Salesforce/blip-image-captioning-base")
except Exception as e:
    print(f"Failed to initialize Hugging Face pipeline: {e}")
    caption_pipeline = None

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

@app.route('/upload', methods=['POST'])
def upload_image():
    global ngrok_process
    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
        caption = "Default caption: Uploaded image"
        if caption_pipeline:
            try:
                image = Image.open(io.BytesIO(image_bytes))
                caption_result = caption_pipeline(image)
                caption = caption_result[0]['generated_text']
            except Exception as e:
                print(f"Failed to generate caption: {e}")

        # Get public URL from ngrok
        port = start_http_server()
        public_url = start_ngrok(port)
        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()
            ngrok_process = None

@app.route('/shutdown', methods=['POST'])
def shutdown():
    global httpd, ngrok_process
    try:
        if httpd:
            httpd.shutdown()
            httpd.server_close()
            print("HTTP server shut down")
        if ngrok_process:
            ngrok_process.terminate()
            ngrok_process = None
            print("Ngrok shut down")
        return jsonify({"status": "success", "message": "Server shutdown"})
    except Exception as e:
        return jsonify({"status": "error", "message": str(e)})

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