app.py

import os
import io
import json
import random
import logging
from PIL import Image, ImageDraw
import numpy as np
import requests
import gradio as gr

# Force CPU execution configurations for PyTorch to prevent memory bloat on HF Spaces
os.environ["OMP_NUM_THREADS"] = "1"
os.environ["MKL_NUM_THREADS"] = "1"

import torch
from torchvision import transforms

# --- Logging Setup ---
logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
logger = logging.getLogger("CompleteProfileAI")

# --- In-Memory ML Model Initialization ---
DEVICE = "cpu"
BIREFNET_MODEL = None

try:
    from transformers import AutoModelForImageSegmentation
    logger.info("Initializing BiRefNet on CPU...")
    # Loading the official lightweight/general BiRefNet weights
    BIREFNET_MODEL = AutoModelForImageSegmentation.from_pretrained(
        "ZhengPeng7/BiRefNet", 
        trust_remote_code=True
    )
    BIREFNET_MODEL.to(DEVICE)
    BIREFNET_MODEL.eval()
    logger.info("BiRefNet successfully loaded and optimized for CPU.")
except Exception as e:
    logger.error(f"Failed to load local BiRefNet model: {e}. Falling back to transparent bypass.")

# --- Helper 1: Programmatic Studio Gradient Backdrops ---
def create_gradient_backdrop(style="neutral_gray", size=(1024, 1024)):
    """Generates beautiful, professional linear gradient canvases directly in-memory."""
    width, height = size
    base = Image.new("RGB", size)
    
    if style == "office_blue":
        color1 = (15, 32, 67)      # Deep Corporate Navy
        color2 = (44, 83, 130)     # Clean Soft Slate Blue
    elif style == "soft_teal":
        color1 = (11, 40, 41)      # Dark Forest Teal
        color2 = (41, 108, 104)    # Warm Modern Muted Teal
    else: # neutral_gray
        color1 = (25, 25, 25)      # Rich Charcoal
        color2 = (85, 85, 85)      # Soft Studio Medium Gray
        
    for y in range(height):
        ratio = y / height
        r = int(color1[0] * (1 - ratio) + color2[0] * ratio)
        g = int(color1 * (1 - ratio) + color2 * ratio)
        b = int(color1 * (1 - ratio) + color2 * ratio)
        # Apply the row color efficiently
        for x in range(width):
            base.putpixel((x, y), (r, g, b))
            
    return base

# --- Helper 2: Programmatic Abstract Fallback Banners ---
def generate_fallback_banner(industry, color_palette_name):
    """Generates a beautiful geometric abstract banner programmatically if the API is offline."""
    size = (1584, 396)
    image = Image.new("RGB", size)
    draw = ImageDraw.Draw(image)
    
    # Established Corporate Brand Palettes
    palettes = {
        "Corporate Blue": ((15, 32, 67), (44, 83, 130), (100, 149, 237)),
        "Creative Teal": ((11, 40, 41), (41, 108, 104), (127, 255, 212)),
        "Tech Slate": ((15, 15, 15), (50, 50, 60), (150, 150, 160)),
        "Creative Amber": ((60, 30, 10), (120, 60, 20), (255, 191, 0)),
    }
    
    colors = palettes.get(color_palette_name, palettes["Corporate Blue"])
    c1, c2, c3 = colors
    
    # Set background gradient
    for y in range(396):
        ratio = y / 396
        r = int(c1[0] * (1 - ratio) + c2[0] * ratio)
        g = int(c1 * (1 - ratio) + c2 * ratio)
        b = int(c1 * (1 - ratio) + c2 * ratio)
        draw.line([(0, y), (1584, y)], fill=(r, g, b))
        
    # Generate abstract overlapping translucent shapes seeded by industry
    random.seed(hash(industry))
    
    for _ in range(12):
        x1 = random.randint(0, 1584)
        y1 = random.randint(0, 396)
        x2 = x1 + random.randint(100, 450)
        y2 = y1 + random.randint(50, 300)
        x3 = x1 + random.randint(-200, 200)
        y3 = y1 + random.randint(-150, 150)
        
        # Overlay translucent polygon on top of the base
        shape_img = Image.new("RGBA", size)
        shape_draw = ImageDraw.Draw(shape_img)
        fill_color = random.choice([c2, c3]) + (random.randint(25, 75),) # RGB + Alpha
        shape_draw.polygon([(x1, y1), (x2, y2), (x3, y3)], fill=fill_color)
        image = Image.alpha_composite(image.convert("RGBA"), shape_img).convert("RGB")
        
    return image

# --- CORE FUNCTION 1: AI Career Journalist (Tab 1) ---
def optimize_linkedin_text(target_role, core_skills, achievement, style):
    api_key = os.environ.get("OPENAI_API_KEY")
    if not api_key:
        logger.warning("OPENAI_API_KEY environment variable not configured.")
        return (
            "OpenAI API Key is missing. Please add it to your Space Secrets in settings.",
            "Please configure your OPENAI_API_KEY setting inside Hugging Face Secrets to activate the copywriter AI.",
            ["No Key Added"]
        )

    # Initialize OpenAI Client (Compatible with v1.0.0+)
    from openai import OpenAI
    client = OpenAI(api_key=api_key)

    system_prompt = """You are the "AI Career Journalist," an elite Executive Recruiter and world-class LinkedIn Copywriter. Your mission is to interview job seekers and translate their messy, unstructured, conversational raw inputs into compelling, high-converting, and keyword-optimized LinkedIn profiles. 

Your copywriting philosophy:
1. Cut the Fluff: Avoid generic corporate corporate-speak. Be concrete.
2. Quantify Impact: Turn passive duties into active, measurable achievements.
3. Keep it Human: Write in a natural, professional first-person tone ("I am...", "I lead...") that sounds like a confident professional, not an LLM.

You must strictly output your response in valid JSON format matching the schema requested. Do not write conversational preambles or postscripts."""

    task_prompt = f"""Transform the following conversational user inputs into a polished LinkedIn Headline, "About" Summary, and Keyword List.

### Few-Shot Example:
- USER INPUTS:
  * Target Role: Junior Software Engineer
  * Core Skills: Python, React, PostgreSQL, Git
  * Major Achievement: Built a campus tutoring app that was used by 300 students to schedule sessions.
  * Working Style: Collaborative, analytical, eager to solve complex logic.
- AI OUTPUT JSON:
{{
  "headline": "Junior Software Engineer | Python & React | Building Impact-Driven Web Solutions",
  "summary": "I am a software engineer focused on building highly functional, user-centric web applications. My passion lies in translating complex logic into clean, performant code.\\n\\nRecently, I developed a campus tutoring scheduler using Python and React, which successfully streamlined session booking for over 300 active student users. I thrive in collaborative environments where continuous learning and analytical problem-solving are valued.\\n\\nSpecialties: Python, JavaScript (React), SQL (PostgreSQL), Git, API Integration, and Agile Methodologies.",
  "extracted_keywords": ["Software Engineering", "Full-Stack Development", "Python", "React.js", "PostgreSQL", "Database Design", "Agile Methodologies"],
  "status": "success",
  "error_message": ""
}}

### Live Task:
- USER INPUTS:
  * Target Role: {target_role}
  * Core Skills: {core_skills}
  * Major Achievement: {achievement}
  * Working Style: {style}

Generate the JSON response following the exact schema shown in the examples. Your output MUST be pure JSON with no markdown wrapping (i.e. no ```json)."""

    try:
        response = client.chat.completions.create(
            model="gpt-4o-mini",
            response_format={ "type": "json_object" }, # Force structured JSON format
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": task_prompt}
            ],
            timeout=30
        )
        content = response.choices[0].message.content
        data = json.loads(content)
        
        if data.get("status") == "error":
            return "", data.get("error_message", "Error compiling data."), []
            
        return data["headline"], data["summary"], data["extracted_keywords"]
    except Exception as e:
        logger.error(f"Error during OpenAI copy generation: {e}")
        return (
            f"{target_role} | {core_skills.split(',')[0] if core_skills else 'Professional'}",
            f"An error occurred while calling the OpenAI service: {e}. Please ensure your API secrets are configured correctly.",
            [s.strip() for s in core_skills.split(",")] if core_skills else []
        )

# --- CORE FUNCTION 2: Studio Headshot background Remover (Tab 2) ---
def process_headshot(image, bg_type, gradient_preset, solid_color):
    if image is None:
        return None
        
    # Step 1: Preprocess size to safeguard against CPU RAM overflow
    max_size = 1024
    w, h = image.size
    if w > max_size or h > max_size:
        ratio = min(max_size / w, max_size / h)
        new_size = (int(w * ratio), int(h * ratio))
        image = image.resize(new_size, Image.Resampling.LANCZOS)
        logger.info(f"Resized input image to safe processing boundaries: {new_size}")
        
    orig_w, orig_h = image.size
    
    # Step 2: Extract Subject using Local BiRefNet Model
    if BIREFNET_MODEL is None:
        # Transparent Bypass Fallback if PyTorch fails to load
        logger.warning("BiRefNet uninitialized. Bypassing background extraction.")
        cut_subject = image.convert("RGBA")
    else:
        try:
            logger.info("Starting BiRefNet background segmentation loop...")
            # Normalize and prepare input tensor
            transform_image = transforms.Compose([
                transforms.Resize((1024, 1024)),
                transforms.ToTensor(),
                transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
            ])
            img_rgb = image.convert("RGB")
            input_tensor = transform_image(img_rgb).unsqueeze(0).to(DEVICE)
            
            with torch.no_grad():
                outputs = BIREFNET_MODEL(input_tensor)
                # Defensively unpack outputs based on tensor structure
                if hasattr(outputs, "logits"):
                    pred = outputs.logits[-1] if isinstance(outputs.logits, list) else outputs.logits
                elif isinstance(outputs, (list, tuple)):
                    pred = outputs[-1]
                else:
                    pred = outputs
                    
                if len(pred.shape) == 4:
                    pred = pred.squeeze(0).squeeze(0)
                elif len(pred.shape) == 3:
                    pred = pred.squeeze(0)
                    
                # Generate alpha mask through sigmoid mapping
                pred = torch.sigmoid(pred).cpu().numpy()
                
            # Resize binary mask back to original bounds
            mask = Image.fromarray((pred * 255).astype(np.uint8)).resize((orig_w, orig_h), Image.Resampling.BILINEAR)
            cut_subject = image.convert("RGBA")
            cut_subject.putalpha(mask)
            logger.info("Successfully extracted headshot foreground subject.")
        except Exception as e:
            logger.error(f"Error during BiRefNet execution: {e}")
            cut_subject = image.convert("RGBA")
            
    # Step 3: Overlay Foreground onto Selected Studio Backdrop
    if bg_type == "Solid Color":
        hex_color = solid_color.lstrip('#')
        rgb_color = tuple(int(hex_color[i:i+2], 16) for i in (0, 2, 4))
        background = Image.new("RGB"