import os
import torch
import boto3
import random
import string
import numpy as np
import logging
import datetime
from fastapi import FastAPI, HTTPException, Request, Response
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, constr, conint
from diffusers import (FluxPipeline, FluxControlNetPipeline, 
                        FluxControlNetModel, FluxImg2ImgPipeline, 
                        FluxInpaintPipeline, CogVideoXImageToVideoPipeline)
from diffusers.utils import load_image
from PIL import Image

# Setup logging
logging.basicConfig(level=logging.INFO, 
                    format='%(asctime)s - %(levelname)s - %(message)s',
                    handlers=[
                        logging.FileHandler("error.txt"),
                        logging.StreamHandler()
                    ])

app = FastAPI()

# Allow CORS for specific origins if needed
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # Update with specific domains as necessary
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

MAX_SEED = np.iinfo(np.int32).max

# AWS S3 Configuration
AWS_ACCESS_KEY_ID = "your-access-key-id"
AWS_SECRET_ACCESS_KEY = "your-secret-access-key"
AWS_REGION = "your-region"
S3_BUCKET_NAME = "your-bucket-name"

# Initialize S3 client
s3_client = boto3.client(
    's3',
    aws_access_key_id=AWS_ACCESS_KEY_ID,
    aws_secret_access_key=AWS_SECRET_ACCESS_KEY,
    region_name=AWS_REGION
)

def log_requests(user_key: str, prompt: str):
    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    log_entry = f"{timestamp}, {user_key}, {prompt}\n"
    with open("key_requests.txt", "a") as log_file:
        log_file.write(log_entry)

# Function to upload image to S3
def upload_image_to_s3(image_path: str, s3_path: str):
    try:
        s3_client.upload_file(image_path, S3_BUCKET_NAME, s3_path)
        return f"https://{S3_BUCKET_NAME}.s3.{AWS_REGION}.amazonaws.com/{s3_path}"
    except Exception as e:
        logging.error(f"Error uploading image to S3: {e}")
        raise HTTPException(status_code=500, detail=f"Image upload failed: {str(e)}")

# Generate a random sequence of 12 numbers and 11 words
def generate_random_sequence():
    random_numbers = ''.join(random.choices(string.digits, k=12))  # 12 random digits
    random_words = ''.join(random.choices(string.ascii_lowercase, k=11))  # 11 random letters
    return f"{random_numbers}_{random_words}"

# Load the default pipeline once globally for efficiency

# Load the default pipeline once globally for efficiency
try:
    flux_pipe = FluxPipeline.from_pretrained("pranavajay/flow", torch_dtype=torch.bfloat16)
    flux_pipe.enable_model_cpu_offload()
    logging.info("FluxPipeline loaded successfully.")
except Exception as e:
    logging.error(f"Failed to load FluxPipeline: {e}")
    raise HTTPException(status_code=500, detail=f"Failed to load the model: {str(e)}")
    
try:
    img_pipe = FluxImg2ImgPipeline.from_pretrained("pranavajay/flow", torch_dtype=torch.bfloat16)
    img_pipe.enable_model_cpu_offload()
    logging.info("FluxImg2ImgPipeline loaded successfully.")
except Exception as e:
    logging.error(f"Failed to load FluxPipeline: {e}")
    raise HTTPException(status_code=500, detail=f"Failed to load the model: {str(e)}")
    
try:
    inpainting_pipe = FluxInpaintPipeline.from_pretrained("pranavajay/flow", torch_dtype=torch.bfloat16)
    inpainting_pipe.enable_model_cpu_offload()
    logging.info("FluxInpaintPipeline loaded successfully.")
except Exception as e:
    logging.error(f"Failed to load FluxInpaintPipeline: {e}")
    raise HTTPException(status_code=500, detail=f"Failed to load the model: {str(e)}")
    
try:
    video = CogVideoXImageToVideoPipeline.from_pretrained(
        "THUDM/CogVideoX-5b-I2V",
        torch_dtype=torch.bfloat16
    )
    video.enable_sequential_cpu_offload()
    video.vae.enable_tiling()
    video.vae.enable_slicing()
    logging.info("CogVideoXImageToVideoPipeline loaded successfully.")
except Exception as e:
    logging.error(f"Failed to load CogVideoXImageToVideoPipeline: {e}")
    raise HTTPException(status_code=500, detail=f"Failed to load the model: {str(e)}")


flux_controlnet_pipe = None



# Rate limiting variables
request_timestamps = defaultdict(list)  # Store timestamps of requests per user key
RATE_LIMIT = 30  # Maximum requests allowed
TIME_WINDOW = 5  # Time window in seconds

# Available LoRA styles and ControlNet adapters
style_lora_mapping = {
    "Uncensored": {"path": "enhanceaiteam/Flux-uncensored", "triggered_word": "nsfw"},
    "Logo": {"path": "Shakker-Labs/FLUX.1-dev-LoRA-Logo-Design", "triggered_word": "logo"},
    "Yarn": {"path": "Shakker-Labs/FLUX.1-dev-LoRA-MiaoKa-Yarn-World", "triggered_word": "mkym this is made of wool"},
    "Anime": {"path": "prithivMLmods/Canopus-LoRA-Flux-Anime", "triggered_word": "anime"},
    "Comic": {"path": "wkplhc/comic", "triggered_word": "comic"}
}

adapter_controlnet_mapping = {
    "Canny": "InstantX/FLUX.1-dev-controlnet-canny",
    "Depth": "Shakker-Labs/FLUX.1-dev-ControlNet-Depth",
    "Pose": "Shakker-Labs/FLUX.1-dev-ControlNet-Pose",
    "Upscale": "jasperai/Flux.1-dev-Controlnet-Upscaler"
}

# Request model for query parameters
class GenerateImageRequest(BaseModel):
    prompt: constr(min_length=1)  # Ensures prompt is not empty
    guidance_scale: float = 7.5
    seed: conint(ge=0, le=MAX_SEED) = 42
    randomize_seed: bool = False
    height: conint(gt=0) = 768
    width: conint(gt=0) = 1360
    control_image_url: str = "https://enhanceai.s3.amazonaws.com/792e2322-77fe-4070-aac4-7fa8d9e29c11_1.png"
    controlnet_conditioning_scale: float = 0.6
    num_inference_steps: conint(gt=0) = 50
    num_images_per_prompt: conint(gt=0, le=5) = 1  # Limit to max 5 images per request
    style: str = None  # Optional LoRA style
    adapter: str = None  # Optional ControlNet adapter
    user_key: str  # API user key

def log_request(key: str, query: str):
    with open("key.txt", "a") as f:
        f.write(f"{datetime.datetime.now()} - Key: {key} - Query: {query}\n")

def apply_lora_style(pipe, style, prompt):
    """ Apply the specified LoRA style to the prompt and load weights. """
    if style in style_lora_mapping:
        lora_path = style_lora_mapping[style]["path"]
        triggered_word = style_lora_mapping[style]["triggered_word"]
        pipe.load_lora_weights(lora_path)
        return f"{triggered_word} {prompt}"  # Add triggered word to prompt
    return prompt

def set_controlnet_adapter(adapter: str, is_inpainting: bool = False):
    """
    Set the ControlNet adapter for the pipeline.
    
    Parameters:
    adapter (str): The key to identify which ControlNet adapter to load.
    is_inpainting (bool, optional): Whether to use the inpainting pipeline. Defaults to False.
    
    Raises:
    ValueError: If the adapter is not found in the adapter_controlnet_mapping.
    """
    global flux_controlnet_pipe

    # Check if the adapter is valid
    if adapter not in adapter_controlnet_mapping:
        raise ValueError(f"Invalid ControlNet adapter: {adapter}")

    # Get the ControlNet model path based on the adapter
    controlnet_model_path = adapter_controlnet_mapping[adapter]
    
    # Load the ControlNet model with the specified torch_dtype
    controlnet = FluxControlNetModel.from_pretrained(controlnet_model_path, torch_dtype=torch.bfloat16)
    
    # Select the appropriate pipeline (inpainting or standard)
    pipeline_cls = FluxControlNetInpaintPipeline if is_inpainting else FluxControlNetPipeline
    
    # Load the pipeline
    flux_controlnet_pipe = pipeline_cls.from_pretrained(
        "pranavajay/flow", controlnet=controlnet, torch_dtype=torch.bfloat16
    )

    # Move the pipeline to the GPU
    flux_controlnet_pipe.to("cuda")

    logging.info(f"ControlNet adapter '{adapter}' loaded successfully.")



        

        

def rate_limit(user_key: str):
    """ Check if the user is exceeding the rate limit. """
    current_time = time.time()

    # Clean up old timestamps
    request_timestamps[user_key] = [t for t in request_timestamps[user_key] if current_time - t < TIME_WINDOW]

    if len(request_timestamps[user_key]) >= RATE_LIMIT:
        logging.info(f"Rate limit exceeded for user_key: {user_key}")
        return False

    # Record the new request timestamp
    request_timestamps[user_key].append(current_time)
    return True

@app.post("/text_to_image/")
async def generate_image(req: GenerateImageRequest):
    seed = req.seed
   if not rate_limit(req.user_key):
        log_requests(req.user_key, req.prompt)  # Log the request when rate limit is exceeded
    
    retries = 3  # Number of retries for transient errors

    for attempt in range(retries):
        try:
            # Check if prompt is None or empty
            if not req.prompt or req.prompt.strip() == "":
                raise ValueError("Prompt cannot be empty.")
            
            original_prompt = req.prompt  # Save the original prompt
            
            # Set ControlNet if adapter is provided
            if req.adapter:
                try:
                    set_controlnet_adapter(req.adapter)
                except Exception as e:
                    logging.error(f"Error setting ControlNet adapter: {e}")
                    raise HTTPException(status_code=400, detail=f"Failed to load ControlNet adapter: {str(e)}")
                    apply_lora_style(flux_controlnet_pipe, req.style, req.prompt)
                    
                
                # Load control image
                try:
                    control_image = load_image(req.control_image_url)
                except Exception as e:
                    logging.error(f"Error loading control image from URL: {e}")
                    raise HTTPException(status_code=400, detail="Invalid control image URL or image could not be loaded.")
                
                # Image generation with ControlNet
                try:
                    if req.randomize_seed:
                        seed = random.randint(0, MAX_SEED)
                    generator = torch.Generator().manual_seed(seed)

                    images = flux_controlnet_pipe(
                        prompt=modified_prompt,
                        guidance_scale=req.guidance_scale,
                        height=req.height,
                        width=req.width,
                        num_inference_steps=req.num_inference_steps,
                        num_images_per_prompt=req.num_images_per_prompt,
                        control_image=control_image,
                        generator=generator,
                        controlnet_conditioning_scale=req.controlnet_conditioning_scale
                    ).images
                except torch.cuda.OutOfMemoryError:
                    logging.error("GPU out of memory error while generating images with ControlNet.")
                    raise HTTPException(status_code=500, detail="GPU overload occurred while generating images. Try reducing the resolution or number of steps.")
                except Exception as e:
                    logging.error(f"Error during image generation with ControlNet: {e}")
                    raise HTTPException(status_code=500, detail=f"Error during image generation: {str(e)}")
            else:
                # Image generation without ControlNet
                try:
                    apply_lora_style(flux_pipe, req.style, req.prompt)
                    if req.randomize_seed:
                        seed = random.randint(0, MAX_SEED)
                    generator = torch.Generator().manual_seed(seed)

                    images = flux_pipe(
                        prompt=modified_prompt,
                        guidance_scale=req.guidance_scale,
                        height=req.height,
                        width=req.width,
                        num_inference_steps=req.num_inference_steps,
                        num_images_per_prompt=req.num_images_per_prompt,
                        generator=generator
                    ).images
                except torch.cuda.OutOfMemoryError:
                    logging.error("GPU out of memory error while generating images without ControlNet.")
                    raise HTTPException(status_code=500, detail="GPU overload occurred while generating images. Try reducing the resolution or number of steps.")
                except Exception as e:
                    logging.error(f"Error during image generation without ControlNet: {e}")
                    raise HTTPException(status_code=500, detail=f"Error during image generation: {str(e)}")

            # Saving images and uploading to S3
            image_urls = []
            for i, img in enumerate(images):
                image_path = f"generated_images/{generate_random_sequence()}.png"
                img.save(image_path)
                image_url = upload_image_to_s3(image_path, image_path)
                image_urls.append(image_url)
                os.remove(image_path)  # Clean up local files after upload
            
            return {"status": "success", "output": image_url, "prompt": original_prompt, "height": req.height, "width": req.width, "scale": req.guidance_scale, "step": step, "sytle": req.sytle, "adapter": req.adapter}

        except Exception as e:
            logging.error(f"Attempt {attempt + 1} failed: {e}")
            if attempt == retries - 1:  # Last attempt
                raise HTTPException(status_code=500, detail=f"Failed to generate image after multiple attempts: {str(e)}")
            continue  # Retry on transient errors
            
            
# Image-to-Image request model
class GenerateImageToImageRequest(BaseModel):
    prompt: str = None  # Prompt can be None
    image: str = "https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo.png"
    strength: float = 0.7
    guidance_scale: float = 7.5
    seed: conint(ge=0, le=MAX_SEED) = 42
    randomize_seed: bool = False
    height: conint(gt=0) = 768
    width: conint(gt=0) = 1360
    control_image_url: str = None  # Optional ControlNet image
    controlnet_conditioning_scale: float = 0.6
    num_inference_steps: conint(gt=0) = 50
    num_images_per_prompt: conint(gt=0, le=5) = 1
    style: str = None  # Optional LoRA style
    adapter: str = None  # Optional ControlNet adapter
    user_key: str  # API user key

@app.post("/image_to_image/")
async def generate_image_to_image(req: GenerateImageToImageRequest):
    seed = req.seed
    original_prompt = req.prompt
    modified_prompt = original_prompt
    
    # Check if user is exceeding rate limit
    if not rate_limit(req.user_key):
        log_requests(req.user_key, req.prompt if req.prompt else "No prompt")
        raise HTTPException(status_code=429, detail="Rate limit exceeded")
   
    retries = 3  # Number of retries for transient errors

    for attempt in range(retries):
        try:
            # Check if prompt is None or empty
            if not req.prompt or req.prompt.strip() == "":
                raise ValueError("Prompt cannot be empty.")
            
            original_prompt = req.prompt  # Save the original prompt
            
            # Set ControlNet if adapter is provided
            if req.adapter:
                try:
                    set_controlnet_adapter(req.adapter)
                except Exception as e:
                    logging.error(f"Error setting ControlNet adapter: {e}")
                    raise HTTPException(status_code=400, detail=f"Failed to load ControlNet adapter: {str(e)}")
                    apply_lora_style(flux_controlnet_pipe, req.style, req.prompt)
                    
                
                # Load control image
                try:
                    control_image = load_image(req.control_image_url)
                except Exception as e:
                    logging.error(f"Error loading control image from URL: {e}")
                    raise HTTPException(status_code=400, detail="Invalid control image URL or image could not be loaded.")
                
                # Image generation with ControlNet
                try:
                    if req.randomize_seed:
                        seed = random.randint(0, MAX_SEED)
                    generator = torch.Generator().manual_seed(seed)
                    
                    images = flux_controlnet_pipe(
                        prompt=modified_prompt,
                        guidance_scale=req.guidance_scale,
                        height=req.height,
                        width=req.width,
                        num_inference_steps=req.num_inference_steps,
                        num_images_per_prompt=req.num_images_per_prompt,
                        control_image=control_image,
                        generator=generator,
                        controlnet_conditioning_scale=req.controlnet_conditioning_scale
                    ).images
                except torch.cuda.OutOfMemoryError:
                    logging.error("GPU out of memory error while generating images with ControlNet.")
                    raise HTTPException(status_code=500, detail="GPU overload occurred while generating images. Try reducing the resolution or number of steps.")
                except Exception as e:
                    logging.error(f"Error during image generation with ControlNet: {e}")
                    raise HTTPException(status_code=500, detail=f"Error during image generation: {str(e)}")
            else:
                # Image generation without ControlNet
                try:
                    apply_lora_style(img_pipe, req.style, req.prompt)
                    if req.randomize_seed:
                        seed = random.randint(0, MAX_SEED)
                    generator = torch.Generator().manual_seed(seed)
                    source = load_image(req.image)
                    images = img_pipe(
                        prompt=modified_prompt,
                        image=source,
                        strength=req.strength,
                        guidance_scale=req.guidance_scale,
                        height=req.height,
                        width=req.width,
                        num_inference_steps=req.num_inference_steps,
                        num_images_per_prompt=req.num_images_per_prompt,
                        generator=generator
                    ).images
                except torch.cuda.OutOfMemoryError:
                    logging.error("GPU out of memory error while generating images without ControlNet.")
                    raise HTTPException(status_code=500, detail="GPU overload occurred while generating images. Try reducing the resolution or number of steps.")
                except Exception as e:
                    logging.error(f"Error during image generation without ControlNet: {e}")
                    raise HTTPException(status_code=500, detail=f"Error during image generation: {str(e)}")

            # Saving images and uploading to S3
            image_urls = []
            for i, img in enumerate(images):
                image_path = f"generated_images/{generate_random_sequence()}.png"
                img.save(image_path)
                image_url = upload_image_to_s3(image_path, image_path)
                image_urls.append(image_url)
                os.remove(image_path)  # Clean up local files after upload
            
            return {"status": "success", "output": image_url, "prompt": original_prompt, "height": req.height, "width": width, "image": req.image, "strength": req.strength, "scale": req.guidance_scale, "step": step, "sytle": req.sytle, "adapter": req.adapter}

        except Exception as e:
            logging.error(f"Attempt {attempt + 1} failed: {e}")
            if attempt == retries - 1:  # Last attempt
                raise HTTPException(status_code=500, detail=f"Failed to generate image after m