import uuid import boto3 import requests import time from typing import Dict, Any from io import BytesIO from app.models.generated_tattoo_history import GeneratedTattooHistory from fastapi import APIRouter, Depends, Request from pydantic import BaseModel, ValidationError from sqlalchemy.ext.asyncio import AsyncSession from app.schemas.store_s3 import BgRemovalRequest from app.db.session import get_db from app.models.bg_remove import BgRemoveImg import urllib.parse from app.services.store_image_s3 import upload_image_to_s3 from app.core.security import validate_token from app.utils.responses import custom_response from app.api.v1.category_api import decrypt_data from app.core.logging import logger from app.core.config import settings from types import SimpleNamespace import random import httpx import asyncio from sqlalchemy import update, select import base64 def now_epoch() -> int: """ Returns the current time in Unix epoch format (integer). """ return int(time.time()) router = APIRouter() def generate_random_public_ip(): """Generate a random public IP address (avoiding private ranges)""" while True: ip = [random.randint(0, 255) for _ in range(4)] # Skip private IP ranges if (ip[0] == 10 or (ip[0] == 172 and 16 <= ip[1] <= 31) or (ip[0] == 192 and ip[1] == 168) or ip[0] == 127): continue return '.'.join(map(str, ip)) def build_response_from_history(history_obj, extra_fields=None): """Builds a consistent response dict from GeneratedTattooHistory or BgRemoveImg""" if not history_obj: return {} response = { "history_id": getattr(history_obj, "history_id", None), "bg_history_id": getattr(history_obj, "bg_history_id", None), # for bg_remove "user_id": getattr(history_obj, "user_id", None) or getattr(history_obj, "user_device_id", None), "tattoo_style_id": getattr(history_obj, "tattoo_style_id", None), "name": getattr(history_obj, "name", None), "search_name": getattr(history_obj, "search_name", None), "image_url": getattr(history_obj, "image_url", None), "bg_url": getattr(history_obj, "bg_url", None), "placement_name": getattr(history_obj, "placement_name", None), "text_content": getattr(history_obj, "text_content", None), "font_family": getattr(history_obj, "font_family", None), "size": getattr(history_obj, "size", None), "alignment": getattr(history_obj, "alignment", None), "color_code": getattr(history_obj, "color", None), "color_name": getattr(history_obj, "color_name", None), "tattoo_type": getattr(history_obj, "tattoo_type", None), "category_type": getattr(history_obj, "category_type", None), "is_delete": getattr(history_obj, "is_delete", None), "is_history": getattr(history_obj, "is_history", None), "is_template": getattr(history_obj, "is_template", None), "created_at": getattr(history_obj, "created_at", None), "updated_at": getattr(history_obj, "updated_at", None), } # Merge in any extra fields (like AI/sketch specific) if extra_fields: response.update(extra_fields) return response # ---------- Background Removal API (POST with encrypted payload) ---------- @router.post("/tattoo/bg_remove") async def background_removal( request: Request, db: AsyncSession = Depends(get_db), token_data: dict = Depends(validate_token) ): """ Removes the background from an image using the Pixian AI API, saves the result to S3, and stores the URL in the database. Supports: 1. Encrypted JSON in raw body 2. Multipart form-data with `dataset` (encrypted JSON) + `file` (image) """ logger.info("Received request for /tattoo/bg_remove endpoint.") try: # Step 1: Authorization check if token_data.get('code') != 200: logger.warning(f"Unauthorized access attempt to /tattoo/bg_remove. Message: {token_data.get('message', 'N/A')}") return custom_response(code=token_data.get('code', 500), message=token_data.get('message', 'Unauthorized')) user_device_id = token_data.get('user_device_id') tattoo_style_id = token_data.get('tattoo_style_id') user_id = token_data.get('user_id') if not user_device_id: logger.warning("User ID or device ID not found in token.") return custom_response(code=401, message="User not authenticated") logger.info(f"Processing background removal request for user_device_id: {user_device_id}") # Step 2: Check content type content_type = request.headers.get("content-type", "") logger.debug(f"Request Content-Type: {content_type}") encrypted_data = None upload_file = None print('content_type===================',content_type) if "multipart/form-data" in content_type: # Multipart case: get dataset + file form = await request.form() encrypted_data = form.get("dataset") upload_file = form.get("file") # starlette.datastructures.UploadFile sketch_file = form.get("sketch") api_type =form.get("api_type") logger.debug("Received multipart form-data request") else: # Raw body case print("==============else") return custom_response(code=400, message="Missing dataset payload") if not encrypted_data and not upload_file and not sketch_file: logger.warning("No encrypted dataset or image found in request.") return custom_response(code=400, message="Missing dataset payload or image") # Step 3: Decrypt the dataset print('encrypted_data============',encrypted_data) if encrypted_data: decrypted_data = decrypt_data(encrypted_data) logger.debug(f"Decrypted dataset: {decrypted_data}") print("Decrypted dataset:", decrypted_data) try: decrypted_json = decrypted_data if decrypted_json: params = BgRemovalRequest.model_validate(decrypted_json) print('params-------------',params) logger.info(f"Successfully validated dataset. Name: {params.name}, Image URL: {params.image_url if params.image_url else 'N/A'}") except Exception as ve: logger.warning(f"Validation failed for background removal request. Errors: {str(ve)}") return custom_response(code=422, message="Validation failed", data=str(ve)) else: print('upload_file',upload_file) params = SimpleNamespace(name=None) print("Upload file:", upload_file) # Step 4: Prepare image for Pixian print('api_type',api_type) pixian_files = None pixian_data = {"test": "true"} if api_type == 'ai': image_url_s3 = None if params.tattoo_type != 'text': print('image is coming--------------') textnews = params.text_content or "" ai_style = params.tattoo_style or "" ai_placement = params.placement_name or "" genrated_image_url = params.image_url or "" is_not_generated = params.is_not_generated or False image_Bytes = params.image_Bytes or None print('is_not_generated',is_not_generated) base = False if not is_not_generated: base = True print('=================ais_not_generatedi') random_ip = generate_random_public_ip() if str(getattr(params, "regenerate", "")).lower() == "true": # if getattr(params, "regenerate", "").lower() == "true": # Add a small random variation to avoid duplicate prompt print('Regenrate request') random_suffix = random.choice(["red", "blue", "green", "yellow", "purple", "orange", "pink","black", "white", "gray", "turquoise", "gold", "silver", "teal", "magenta", "cyan", "lavender", "bronze", "crimson"]) textnews = f"{textnews} {random_suffix}" prompt = f"""Create a large, bold, modern tattoo design of a {textnews} in an {ai_style} style. The design should be highly detailed, visually striking, and dynamic, with clean, sharp lines, dramatic shapes, and strong contrast. Include subtle textures, elegant shading, and balanced composition that draws the eye. The tattoo must be completely standalone on a plain white background, floating in space, with no body, skin, or additional background elements. Color: blackwork or muted ink tones as appropriate. Add creative abstract flourishes and modern design elements to make the tattoo powerful, bold, and highly eye-catching.""" print('Prompt:',prompt) # Set headers headers = { "Accept": "application/json", "Content-Type": "application/json", "x-client-ip": random_ip } # Set request body payload = { "prompt": prompt, "style": ai_style, "user_input": textnews } print('payload',payload) # Encode prompt for URL encoded_prompt = urllib.parse.quote(prompt) # Construct Pollinations URL image_url = "https://internal.users.n8n.cloud/webhook/ai_image_generator" logger.info(f"AI Tattoo Generated URL: {image_url}") # 🔥 Optionally download image and upload to S3 try: MAX_RETRIES = 5 RETRY_DELAY = 2 # 🔥 Retry loop ai_response = None async with httpx.AsyncClient() as client: for attempt in range(1, MAX_RETRIES + 1): try: ai_response = await client.post(image_url,json=payload,headers=headers, timeout=60) result = ai_response.json() image_url = result['result'][0]['url'] print('AI Generated Image URL:', image_url) image_response = requests.get(image_url) print('image_response',image_response) image_response.raise_for_status() image_Bytes = image_response.content break except Exception as e: logger.warning(f"Attempt {attempt}/{MAX_RETRIES} failed: {e}") if attempt < MAX_RETRIES: await asyncio.sleep(RETRY_DELAY * attempt) else: logger.error("AI tattoo generation failed after retries") return custom_response(code=500, message="AI tattoo generation failed after retries") except Exception as e: logger.error(f"AI tattoo generation failed: {e}") return custom_response(code=500, message="AI tattoo generation failed") print('ai_response---TYPE----',type(image_Bytes)) if base: image_bytes = BytesIO(image_Bytes) else: image_data = base64.b64decode(image_Bytes) image_bytes = BytesIO(image_data) file_name = f"tattoos/{user_device_id}/{uuid.uuid4()}.png" image_url_s3 = upload_image_to_s3(image_bytes, file_name) if not image_url_s3: return custom_response(code=500, message="Failed to store AI tattoo image") new_history = GeneratedTattooHistory( user_id=user_id, user_device_id=user_device_id, tattoo_style_id=getattr(params, "tattoo_style_id", None), name=getattr(params, "tattoo_style", None), search_name=getattr(params, "search_name", None), image_url=image_url_s3, placement_name=getattr(params, "placement_name", None), text_content=getattr(params, "text_content", None), font_family=getattr(params, "font_family", None), size=getattr(params, "size", None), alignment=getattr(params, "alignment", None), color=getattr(params, "color_code", None), color_name=getattr(params, "color_name", None), tattoo_type=getattr(params, "tattoo_type", None), category_type=getattr(params, "category_type", None), is_delete=0, is_history=1, is_template=1 if getattr(params, "is_template", None) else 0, created_at=now_epoch(), updated_at=now_epoch() ) db.add(new_history) await db.commit() await db.refresh(new_history) response_data = build_response_from_history(new_history) elif api_type == 'sketch': print('=================sketch') logger.info("Using uploaded file for sketch") file_bytes = await sketch_file.read() pixian_files = {'image': (sketch_file.filename, file_bytes, sketch_file.content_type or "image/png")} image_bytes = BytesIO(file_bytes) file_name = f"tattoos/{user_device_id}/{uuid.uuid4()}.png" image_url = upload_image_to_s3(image_bytes, file_name) new_sktech = GeneratedTattooHistory( user_id=user_id, user_device_id=user_device_id, tattoo_style_id=None, image_url=image_url, category_type='sketch', is_history=True, created_at=now_epoch(), updated_at=now_epoch() ) db.add(new_sktech) await db.commit() await db.refresh(new_sktech) response_data = build_response_from_history(new_sktech) elif api_type == 'bg_remove': print('=================bg_remove') if upload_file: print('==============into this') # If file uploaded, send file logger.info("Using uploaded file for background removal") file_bytes = await upload_file.read() pixian_files = {'image': (upload_file.filename, file_bytes, upload_file.content_type or "image/png")} elif params.image_url: # If only image_url present print('==================',params.image_url) history_id = params.history_id logger.info(f"Using image URL for background removal: {params.image_url}") pixian_data["image.url"] = params.image_url else: logger.warning("No file or image_url provided") async with httpx.AsyncClient() as client: pixian_response = await client.post( "https://api.pixian.ai/api/v2/remove-background", data=pixian_data, files=pixian_files, auth=(settings.BACKGROUND_API_KEY, settings.BACKGROUND_API_SECRET), ) if pixian_response.status_code != httpx.codes.ok: logger.error(f"Pixian AI API error: {pixian_response.status_code}, {pixian_response.text}") return custom_response(code=pixian_response.status_code, message=f"Background removal failed: {pixian_response.text}") # Step 6: Upload result to S3 image_bytes = BytesIO(pixian_response.content) file_name = f"tattoos/{user_device_id}/{uuid.uuid4()}.png" image_url = upload_image_to_s3(image_bytes, file_name) logger.info(f"Uploaded processed image to S3: {image_url}") # Step 7: Save in DB # new_tattoo = BgRemoveImg( # user_device_id=user_device_id, # tattoo_style_id=None, # name=params.name, # image_url=image_url, # is_history=False, # created_at=now_epoch(), # updated_at=now_epoch() # ) if not upload_file: stmt_update = ( update(GeneratedTattooHistory) .where(GeneratedTattooHistory.history_id == history_id) .values( bg_url=image_url, updated_at=now_epoch() ) ) await db.execute(stmt_update) await db.commit() # Step 2: Fetch updated row stmt_select = select(GeneratedTattooHistory).where(GeneratedTattooHistory.history_id == history_id) result = await db.execute(stmt_select) updated_tattoo = result.scalar_one_or_none() print('updated_tattoo',updated_tattoo) response_data = build_response_from_history(updated_tattoo) else: response_data = {"bg_url": image_url} return custom_response(200, data=response_data, message="Image processed, saved to S3, and URL stored successfully") except Exception as e: await db.rollback() logger.error(f"Error in background_removal: {str(e)}", exc_info=True) return custom_response(code=500, message=f"Internal server error: {str(e)}")