import os
import time
from os import path
import tempfile
import uuid
import base64
import mimetypes
import json
import io
import random
import string

import torch
from PIL import Image

from safetensors.torch import load_file
from huggingface_hub import hf_hub_download

# Diffusers 관련 라이브러리
import gradio as gr
from diffusers import FluxPipeline

# Google GenAI 라이브러리
from google import genai
from google.genai import types

#######################################
# 0. 환경설정
#######################################

BASE_DIR = path.dirname(path.abspath(__file__)) if "__file__" in globals() else os.getcwd()
CACHE_PATH = path.join(BASE_DIR, "models")

os.environ["TRANSFORMERS_CACHE"] = CACHE_PATH
os.environ["HF_HUB_CACHE"] = CACHE_PATH
os.environ["HF_HOME"] = CACHE_PATH

# 간단한 타이머 클래스
class timer:
    def __init__(self, method_name="timed process"):
        self.method = method_name
    def __enter__(self):
        self.start = time.time()
        print(f"{self.method} starts")
    def __exit__(self, exc_type, exc_val, exc_tb):
        end = time.time()
        print(f"{self.method} took {str(round(end - self.start, 2))}s")

#######################################
# 1. FLUX 파이프라인 로드
#######################################

if not path.exists(CACHE_PATH):
    os.makedirs(CACHE_PATH, exist_ok=True)

pipe = FluxPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    torch_dtype=torch.bfloat16
)

lora_path = hf_hub_download("ByteDance/Hyper-SD", "Hyper-FLUX.1-dev-8steps-lora.safetensors")
pipe.load_lora_weights(lora_path)
pipe.fuse_lora(lora_scale=0.125)

pipe.to(device="cuda", dtype=torch.bfloat16)

#######################################
# 2. Google GenAI를 통한 이미지 내 텍스트 변환 함수
#######################################

def save_binary_file(file_name, data):
    """Google GenAI에서 응답받은 이진 데이터를 이미지 파일로 저장"""
    with open(file_name, "wb") as f:
        f.write(data)

def generate_by_google_genai(text, file_name, model="gemini-2.0-flash-exp"):
    """
    Google GenAI(gemini) 모델을 통해 이미지/텍스트를 생성하거나 변환.
    - text: 변경할 텍스트나 명령어 등 프롬프트
    - file_name: 원본 이미지(예: .png) 경로
    - model: 사용할 gemini 모델 이름
    """
    # (1) 환경 변수에서 API 키 가져오기 (필수)
    api_key = os.getenv("GAPI_TOKEN", None)
    if not api_key:
        raise ValueError(
            "GAPI_TOKEN 환경 변수가 설정되지 않았습니다. "
            "Google GenAI API를 사용하기 위해서는 GAPI_TOKEN이 필요합니다."
        )

    # (2) Google Client 초기화
    client = genai.Client(api_key=api_key)
    
    # (3) 이미지 업로드
    files = [client.files.upload(file=file_name)]
    
    # (4) gemini에 전달할 Content 준비 (이미지 + 프롬프트)
    contents = [
        types.Content(
            role="user",
            parts=[
                types.Part.from_uri(
                    file_uri=files[0].uri,
                    mime_type=files[0].mime_type,
                ),
                types.Part.from_text(text=text),
            ],
        ),
    ]

    # (5) 생성/변환 설정
    generate_content_config = types.GenerateContentConfig(
        temperature=1,
        top_p=0.95,
        top_k=40,
        max_output_tokens=8192,
        response_modalities=["image", "text"],
        response_mime_type="text/plain",
    )

    text_response = ""
    image_path = None

    # 임시 파일로 이미지 받을 준비
    with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp:
        temp_path = tmp.name
        # 응답 스트림 처리
        for chunk in client.models.generate_content_stream(
            model=model,
            contents=contents,
            config=generate_content_config,
        ):
            if not chunk.candidates or not chunk.candidates[0].content or not chunk.candidates[0].content.parts:
                continue
            candidate = chunk.candidates[0].content.parts[0]

            # inline_data가 있으면 이미지 응답
            if candidate.inline_data:
                save_binary_file(temp_path, candidate.inline_data.data)
                print(f"File of mime type {candidate.inline_data.mime_type} saved to: {temp_path}")
                image_path = temp_path
                break
            else:
                # 이미지 없이 텍스트만 반환되는 경우
                text_response += chunk.text + "\n"
    
    # 업로드한 File 객체 제거
    del files

    return image_path, text_response

#######################################
# 3. 임의 알파벳 생성용 함수
#######################################

def generate_random_letters(length: int) -> str:
    """
    length 길이의 임의의 알파벳 문자열을 생성해 반환.
    대문자+소문자(a-z, A-Z) 범위에서 랜덤 뽑기.
    """
    letters = string.ascii_uppercase + string.ascii_lowercase
    return "".join(random.choice(letters) for _ in range(length))

def fill_text_input_with_random(new_text: str) -> str:
    """
    "새로 바꿀 텍스트"의 길이를 세어, 그 길이만큼 임의의 알파벳 문자열 생성 후 반환.
    이 반환값이 곧 "이미지 안에 들어갈 텍스트" UI(=text_input)에 표시됨.
    """
    length = len(new_text)
    random_letters = generate_random_letters(length)
    return random_letters

#######################################
# 4. Gradio 함수 
#######################################

def generate_initial_image(prompt, text, height, width, steps, scale, seed):
    """
    FLUX 파이프라인을 사용해 '텍스트가 포함된 이미지를' 먼저 생성.
    - prompt 내 <text>를 text로 치환
    - <text>가 없다면 "with clear readable text that says '<text>'"를 자동 붙임
    """
    if "<text>" in prompt:
        combined_prompt = prompt.replace("<text>", text)
    else:
        combined_prompt = f"{prompt} with clear readable text that says '{text}'"

    print(f"[DEBUG] Final combined_prompt: {combined_prompt}")

    with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16), timer("inference"):
        result = pipe(
            prompt=[combined_prompt],
            generator=torch.Generator().manual_seed(int(seed)),
            num_inference_steps=int(steps),
            guidance_scale=float(scale),
            height=int(height),
            width=int(width),
            max_sequence_length=256
        ).images[0]

    return result

def change_text_in_image(original_image, new_text):
    """
    Google GenAI의 gemini 모델을 통해,
    업로드된 이미지 내부의 문구를 `new_text`로 변경해주는 함수.
    (여기서는 new_text가 이미 임의의 알파벳으로 채워진 상태가 됨)
    """
    try:
        with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp:
            original_path = tmp.name
            original_image.save(original_path)

        image_path, text_response = generate_by_google_genai(
            text=f"Change the text in this image to: '{new_text}'",
            file_name=original_path
        )

        if image_path:
            with open(image_path, "rb") as f:
                image_data = f.read()
            modified_img = Image.open(io.BytesIO(image_data))
            return modified_img, ""
        else:
            return None, text_response

    except Exception as e:
        raise gr.Error(f"Error: {e}")

#######################################
# 5. Gradio 인터페이스 구성
#######################################

with gr.Blocks(title="Flux + Google GenAI + Random Text Replacement") as demo:
    gr.Markdown(
        """
        # Flux 기반 이미지 생성 + Google GenAI + 임의 알파벳 치환

        **기능 요약**  
        1) "새로 바꿀 텍스트"에 입력된 문자열 길이를 측정  
        2) 그 길이만큼 **무작위 알파벳**(대소문자)으로 구성된 문자열을 만들어  
        3) "이미지 안에 들어갈 텍스트" 칸에 자동으로 입력 (즉, 실제 바뀔 텍스트는 완전히 랜덤)  
        4) 그 후 Google GenAI(gemini)로 이미지 문자를 교체

        ---
        """
    )

    with gr.Row():
        with gr.Column():
            gr.Markdown("## 1) Step 1: FLUX로 텍스트 포함 이미지 생성")
            prompt_input = gr.Textbox(
                lines=3,
                label="이미지 장면/배경 Prompt (use `<text>` placeholder if you like)",
                placeholder="e.g. A white cat with speech bubble says <text>"
            )
            text_input = gr.Textbox(
                lines=1,
                label="이미지 안에 들어갈 텍스트",
                placeholder="e.g. Hello or 안녕"
            )
            with gr.Accordion("고급 설정 (확장)", open=False):
                height = gr.Slider(label="Height", minimum=256, maximum=1152, step=64, value=512)
                width = gr.Slider(label="Width", minimum=256, maximum=1152, step=64, value=512)
                steps = gr.Slider(label="Inference Steps", minimum=6, maximum=25, step=1, value=8)
                scale = gr.Slider(label="Guidance Scale", minimum=0.0, maximum=10.0, step=0.5, value=3.5)
                seed = gr.Number(label="Seed (reproducibility)", value=1234, precision=0)

            generate_btn = gr.Button("Generate Base Image", variant="primary")
            generated_image = gr.Image(label="Generated Image (with text)", type="pil")

        with gr.Column():
            gr.Markdown(
                """
                ## 2) Step 2: "새로 바꿀 텍스트" 길이를 기준으로  
                임의 알파벳을 "이미지 안에 들어갈 텍스트"로 자동 설정 → 이미지 교체
                """
            )
            new_text_input = gr.Textbox(
                label="새로 바꿀 텍스트 (입력된 길이를 기준으로 무작위 알파벳 생성)",
                placeholder="예) Hello123 (길이 8이면, 8글자 랜덤)"
            )

            # (A) 먼저 "새로 바꿀 텍스트" 글자수를 세어, 무작위 알파벳으로 text_input 채우기
            # (B) 그런 뒤 change_text_in_image() 실행
            modify_btn = gr.Button("Generate Random Letters & Change Text in Image", variant="secondary")

            # 최종 결과
            output_img = gr.Image(label="Modified Image", type="pil")
            output_txt = gr.Textbox(label="(If only text returned)")

    ###########################
    # 1) Step1 버튼 액션
    ###########################
    generate_btn.click(
        fn=generate_initial_image,
        inputs=[prompt_input, text_input, height, width, steps, scale, seed],
        outputs=[generated_image]
    )

    ###########################
    # 2) Step2 버튼 액션
    ###########################
    # (A) 무작위 알파벳 생성해 text_input에 반영
    chain = modify_btn.click(
        fn=fill_text_input_with_random,
        inputs=[new_text_input],
        outputs=[text_input]
    )
    # (B) 방금 업데이트된 text_input 값을 이용해, 이미지 텍스트 교체
    chain.then(
        fn=change_text_in_image,
        inputs=[generated_image, text_input],
        outputs=[output_img, output_txt]
    )

demo.launch(max_threads=20)