import os
import threading
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.responses import  FileResponse
from fastapi.staticfiles import StaticFiles
from huggingface_hub import InferenceClient
from concurrent.futures import ThreadPoolExecutor
from fastapi.responses import HTMLResponse
from groq import Groq
import uvicorn
import base64
from io import BytesIO
import os
import requests

API_URL = "https://api-inference.huggingface.co/models/Falconsai/nsfw_image_detection"
headers = {"Authorization": f"Bearer {os.getenv('HF_TOKEN')}"}

def is_image_safe(data):
    try:
        response = requests.post(API_URL, headers=headers, data=data)
        response.raise_for_status()  # Raise an error for HTTP issues
        result = response.json()
        print(result)

        # Look for the 'nsfw' label and extract its score
        nsfw_score = next((item["score"] for item in result if item["label"] == "nsfw"), None)

        # Ensure the 'nsfw' label exists in the response
        if nsfw_score is not None:
            return nsfw_score < 0.5  # Return True if the NSFW score is less than 0.5
        else:
            print("Error: 'nsfw' label not found in response.")
            return False
    except requests.exceptions.RequestException as e:
        print(f"Error: Failed to process the request. {e}")
        return False


app = FastAPI()

# Serve static files (HTML, CSS, JS)
app.mount("/static", StaticFiles(directory="static"), name="static")

# Initialize a single Groq client and inference client
client = Groq()
image_gen_client = InferenceClient("black-forest-labs/FLUX.1-schnell")
executor = ThreadPoolExecutor()

@app.get("/", response_class=FileResponse)
async def get():
    return FileResponse("static/index.html")


@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept()

    def generate_text(input_text):
        chat_completion = client.chat.completions.create(
            messages=[
                {"role": "system", "content": "You are an assistant that generates gratitude journal entries. Focus solely on expressing gratitude in a concise, meaningful tone. Do not include introductory or concluding statements. Avoid unrelated topics or personal opinions. Keep the entries simple, neutral, and directly reflective of gratitude."},
                {"role": "user", "content": input_text},
            ],
            model="llama3-8b-8192",
        )
        print(chat_completion.choices[0].message.content)
        return chat_completion.choices[0].message.content
    
    def generate_image_prompt(input_text):
        chat_completion = client.chat.completions.create(
            messages=[
                {"role": "system", "content": "You are an assistant that generates image prompts reflecting gratitude journal entries. Create concise, vivid descriptions that visually represent gratitude and positive moments. Avoid adding any introductory or concluding statements. Keep the prompts simple, neutral, and focused solely on the imagery."},
                {"role": "user", "content": input_text},
            ],
            model="llama3-8b-8192",
        )
        print(chat_completion.choices[0].message.content)
        return chat_completion.choices[0].message.content

    def analyze_image(image_data):
        completion = client.chat.completions.create(
            model="llama-3.2-11b-vision-preview",
            messages=[
                {
                    "role": "user",
                    "content": [
                        {"type": "text", "text": "Describe the contents of this image for a gratitude journal."},
                        {"type": "image", "image": image_data},
                    ],
                }
            ],
            temperature=1,
            max_tokens=1024,
            top_p=1,
            stream=False,
            stop=None,
        )
        print(completion.choices[0].message.content)
        return completion.choices[0].message.content

    def handle_text_to_speech(audio_data):
        transcription = client.audio.transcriptions.create(
            file=("audio.mp3", audio_data),
            model="whisper-large-v3-turbo",
            response_format="json",
            language="en",
        )
        return transcription.text

    def generate_image(prompt):
        print(prompt)
        return image_gen_client.text_to_image(prompt)

    try:
        while True:
            data = await websocket.receive_json()
            input_text = data.get("text", "")
            input_images = data.get("images", [])
            input_audio = data.get("audio", None)

            journal_text = input_text
            tasks = []

            # Process images
            for image in input_images:
                def process_image(image=image):
                    image_description = analyze_image(image)
                    nonlocal journal_text
                    journal_text += f" {image_description}"  
                tasks.append(executor.submit(process_image))

            # Process audio
            if input_audio:
                def process_audio():
                    audio_transcription = handle_text_to_speech(input_audio)
                    nonlocal journal_text
                    journal_text += f" {audio_transcription}"  
                tasks.append(executor.submit(process_audio))

            # Wait for all tasks to complete
            for task in tasks:
                task.result()

            # Generate journal entry
            journal_output = generate_text(journal_text)
            journal_image_prompt= generate_image_prompt(journal_text)
            await websocket.send_json({"type": "journal", "content": journal_output})
            await websocket.send_json({"type": "info", "content": "Gratitude-themed image is being generated"})
            image_prompt = f"Generate a gratitude-themed artistic image based on this journal entry: {journal_image_prompt}"
            generated_image = generate_image(image_prompt)  # This returns a PIL.Image object
            print("image generated")
            


            # Convert PIL.Image to Base64
            buffered = BytesIO()
            generated_image.save(buffered, format="PNG")  # Save as PNG or any format you prefer
            img_str=""
            if is_image_safe(buffered.getvalue()):
                img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")  # Base64 encode the image

            # Send the image as a base64 string via WebSocket
            await websocket.send_json({"type": "image", "image": img_str})

            

    except WebSocketDisconnect:
        await websocket.close()