import gradio as gr
import torch
import random
import numpy as np
import datetime

# 履歴保存
from huggingface_hub import HfApi
from huggingface_hub import login
from huggingface_hub import Repository

import os
# HF_TOKEN 環境変数からトークンを明示的に読み込む
hf_token_value = os.getenv("HF_TOKEN")

if hf_token_value:
    api = HfApi(token=hf_token_value)
    print("token ok.")
else:
    # トークンが設定されていない場合の警告と代替処理
    print("HF_TOKEN error")
    api = HfApi() # トークンなしで初期化

# 画像をアップロードするリポジトリID
HF_REPO_ID = "cocoat/images"

# Space内で画像を保存するディレクトリ
SPACE_IMAGE_DIR = "generated_images"
os.makedirs(SPACE_IMAGE_DIR, exist_ok=True)

# Spaceのリポジトリを初期化
# Gradio Spaceでは、カレントディレクトリがSpaceのリポジトリのルートになります。
# HF_SPACE_ID 環境変数には "ユーザー名/スペース名" が入っています。
space_repo_id = "cocoat/Re.cocoamixXL3"
repo = None # repoオブジェクトを初期化
if space_repo_id:
    try:
        repo = Repository(local_dir=".")
        print(f"Successfully initialized Repository object for Space: {space_repo_id}")
    except Exception as e:
        print(f"Failed to initialize Repository object for Space ({space_repo_id}): {e}")
else:
    print("HF_SPACE_ID environment variable not found. Cannot operate on Space repository.")

# 履歴ファイルを定義
HISTORY_FILE = "history/generation_history_coamixXL3.txt"

# 履歴をロードする関数
import os
import requests
def load_history():
    history_data = []
    hf_raw_file_url = f"https://huggingface.co/datasets/{HF_REPO_ID}/raw/main/{HISTORY_FILE}"
    headers = {}
    if hf_token_value:
        headers["Authorization"] = f"Bearer {hf_token_value}"

    try:
        response = requests.get(hf_raw_file_url, headers=headers)
        response.raise_for_status()

        loaded_hub_paths = set() # 重複ロードを防ぐため

        for line in response.text.splitlines():
            parts = line.strip().split("|||")
            if len(parts) == 2:
                image_path_in_repo, caption = parts
                
                filename_from_hub_path = os.path.basename(image_path_in_repo)
                space_local_image_path = os.path.join(SPACE_IMAGE_DIR, filename_from_hub_path)
                
                # Space内にファイルが存在するか確認
                if os.path.exists(space_local_image_path):
                    # 存在するなら Space のパスを使用
                    history_data.append((image_path_in_repo, caption, space_local_image_path))
                    loaded_hub_paths.add(image_path_in_repo)
                else:
                    print(f"Warning: Space image not found for Hub path: {image_path_in_repo}. Skipping for display.")
        print(f"History loaded from Hub and matched with Space images: {len(history_data)} entries.")
    except requests.exceptions.RequestException as e:
        print(f"Error loading history from Hub via raw URL: {e}. Starting with empty history.")
    except Exception as e:
        print(f"An unexpected error occurred while parsing history: {e}. Starting with empty history.")

    return history_data[:10]


# 履歴を初期化時にロード (修正された load_history を使用)
history = load_history()


# 履歴を初期化時にロード
history = load_history()

from PIL import Image
from diffusers import (
    StableDiffusionXLPipeline,
    EulerAncestralDiscreteScheduler,
    DPMSolverMultistepScheduler
)
from huggingface_hub import hf_hub_download, HfApi

# デバイスと型の設定
device = "cuda" if torch.cuda.is_available() else "cpu"
torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
MAX_SEED = np.iinfo(np.int32).max
MAX_SIZE = 2048

# モデルファイルのダウンロード
model_path = hf_hub_download(
    repo_id="cocoat/cocoamix",
    filename="recocoamixXL3_coamixXL3.safetensors"
)

# パイプライン構築
pipe = StableDiffusionXLPipeline.from_single_file(
    model_path,
    torch_dtype=torch_dtype,
    use_safetensors=True
).to(device)

# スケジューラ設定
euler_scheduler = EulerAncestralDiscreteScheduler.from_config(
    pipe.scheduler.config,
    use_karras_sigmas=True
)
dpm_scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
pipe.scheduler = euler_scheduler

def upload_image_to_hub(image_pil, prompt_text):
    # ファイル名を生成（タイムスタンプとプロンプトの一部）
    timestamp = datetime.datetime.now().strftime("%Y%m%d%H%M%S")
    # プロンプトから安全なファイル名の一部を生成
    # safe_prompt = "".join(c for c in prompt_text if c.isalnum() or c in (' ', '.', '_')).replace(' ', '_')[:30]
    filename = f"image_{timestamp}.png"
    filepath = f"temp_{filename}"
    image_pil.save(filepath)

    # Hubにアップロード
    try:
        # リポジトリ内にディレクトリを作成する場合は path_in_repo を使う
        path_in_repo = f"generated_images/{filename}"
        upload_info = api.upload_file(
            path_or_fileobj=filepath,
            path_in_repo=path_in_repo,
            repo_id=HF_REPO_ID,
            repo_type="dataset", # または "space", "model" など、目的のリポジトリタイプ
            # 通常、画像保存には "dataset" タイプのリポジトリが適しています
        )
               # アップロードされたファイルのURLを構築する
        uploaded_file_url = f"https://huggingface.co/datasets/{HF_REPO_ID}/resolve/main/{path_in_repo}"
        print(f"Uploaded {filepath} to {uploaded_file_url}")
        return uploaded_file_url
    except Exception as e:
        print(f"Error uploading image to Hub: {e}")
        return None
    finally:
        # 一時ファイルを削除
        if os.path.exists(filepath):
            os.remove(filepath)
            print(f"一時画像ファイル {filepath} を削除しました。")

def make_html_table(caption):
    formatted_caption = caption.replace("|-|", "\n")
    rows = formatted_caption.split("\n")
    html = '<table style="width:100%;border-collapse:collapse;background:#fffaf1;color:#000">'
    for row in rows:
        if ": " in row:
            key, val = row.split(": ", 1)
            html += (
#                f'{key}: {val}\n'
                f'<tr><th style="text-align:left;border:1px solid #ddd;padding:4px;">{key}</th>'
                f'<td style="border:1px solid #ddd;padding:4px;">{val}</td></tr>'
            )
    html += '</table>'
    return html

def create_dummy_image(width=512, height=512, alpha=0):
    return Image.new("RGBA", (width, height), (0, 0, 0, alpha))

def update_history_tables_on_select(evt: gr.SelectData):
    if evt.index is not None and 0 <= evt.index < len(history):
        selected_caption = history[evt.index][1]
        return make_html_table(selected_caption)
    return ""

def update_history():
    tables_html = "".join(
        f'<div style="margin-bottom:12px">{make_html_table(item[1])}</div>'
        for item in history
    )
    return tables_html

def infer(prompt, neg, seed, rand, w, h, cfg, steps, scheduler_type,
          progress=gr.Progress(track_tqdm=True)):
    if rand:
        seed = random.randint(0, MAX_SEED)
    generator = torch.Generator(device=device).manual_seed(seed)

    pipe.scheduler = euler_scheduler if scheduler_type == "Euler Ancestral" else dpm_scheduler
    pipe.scheduler.set_timesteps(steps)

    def _callback(pipeline, step_idx, timestep, callback_kwargs):
        progress(step_idx / steps, desc=f"Step {step_idx}/{steps}")
        return callback_kwargs

    output = pipe(
        prompt=prompt,
        negative_prompt=neg or None,
        guidance_scale=cfg,
        num_inference_steps=steps,
        width=w,
        height=h,
        generator=generator,
        callback_on_step_end=_callback
    )
    img = output.images[0]

    caption_text = (
        f"Prompt: {prompt}\n"
        f"Negative: {neg or 'None'}\n"
        f"Seed: {seed}\n"
        f"Size: {w}×{h}\n"
        f"CFG: {cfg}\n"
        f"Steps: {steps}\n"
        f"Scheduler: {scheduler_type}"
    )

    caption_text_for_history = caption_text.replace("\n", "|-|").strip()
    
    # 画像をHubにアップロードし、そのURLを取得
    uploaded_image_url = upload_image_to_hub(img, prompt)

    # Space内に画像を保存し、Gitリポジリにコミット・プッシュ
    timestamp = datetime.datetime.now().strftime("%Y%m%d%H%M%S")
    space_image_filename = f"space_image_{timestamp}.png"
    space_image_filepath = os.path.join(SPACE_IMAGE_DIR, space_image_filename)
    
    saved_to_space_git = False
    try:
        img.save(space_image_filepath)
        print(f"Image saved to Space local filesystem: {space_image_filepath}")
        
        # SpaceのGitリポジトリにコミット・プッシュ
        if repo:
            try:
                # Git LFS で追跡するように .gitattributes を確認/追加
                gitattributes_path = os.path.join(SPACE_IMAGE_DIR, ".gitattributes")
                if not os.path.exists(gitattributes_path):
                    with open(gitattributes_path, "w") as f:
                        f.write("*.png filter=lfs diff=lfs merge=lfs -text\n")
                    repo.git_add(gitattributes_path)
                    print(f"Added .gitattributes for LFS tracking in {SPACE_IMAGE_DIR}.")

                repo.git_add(space_image_filepath)
                repo.git_commit(f"Add generated image {space_image_filename}")
                repo.git_push() # これでHubにプッシュされる
                print(f"Image {space_image_filename} committed and pushed to Space repository.")
                saved_to_space_git = True
            except Exception as e:
                print(f"Error committing/pushing image to Space repository: {e}")
        else:
            print("Space repository object not initialized, skipping Git operations for Space.")
            
    except Exception as e:
        print(f"Error saving image to Space local filesystem: {e}")
        saved_to_space_git = False # エラー時にはFalseに設定を保証
    
    # 履歴を更新
    global history
        # Hubへのアップロードが成功し、かつSpaceへのGit保存も成功した場合のみ履歴に追加
    if uploaded_image_url and saved_to_space_git:
        path_in_repo_for_history = uploaded_image_url.split(f"huggingface.co/datasets/{HF_REPO_ID}/resolve/main/")[1]
        history.insert(0, (path_in_repo_for_history, caption_text_for_history, space_image_filepath))
    else:
        print(f"Skipping history update due to failed Hub upload ({uploaded_image_url is None}) or Space Git save ({not saved_to_space_git}).")

    history_max_items = 10
    if len(history) > history_max_items:
        # Space内の最も古い画像を削除
        if history[-1][2] and os.path.exists(history[-1][2]): 
            oldest_space_image_path = history[-1][2] 
            try:
                if repo: # repoオブジェクトが初期化されていればGit操作
                    repo.git_rm(oldest_space_image_path) # Gitからファイルを削除
                    repo.git_commit(f"Remove oldest image {os.path.basename(oldest_space_image_path)}")
                    repo.git_push()
                    print(f"Deleted oldest image from Space repository: {oldest_space_image_path}")
                else: # repoオブジェクトがない場合はローカルファイルのみ削除
                    os.remove(oldest_space_image_path) 
                    print(f"Deleted oldest image from Space local filesystem (no Git): {oldest_space_image_path}")
            except Exception as e:
                print(f"Error deleting oldest image from Space (local/Git): {e}")
        history.pop()

    # 履歴ファイルを更新し、Hubにアップロードする
    temp_history_filepath = "temp_history.txt"
    with open(temp_history_filepath, "w", encoding="utf-8") as f:
        for img_path_in_repo, cap_text, _ in history:
            f.write(f"{img_path_in_repo}|||{cap_text}\n")

    try:
        api.upload_file(
            path_or_fileobj=temp_history_filepath,
            path_in_repo=HISTORY_FILE,
            repo_id=HF_REPO_ID,
            repo_type="dataset",
        )
        print(f"History file '{HISTORY_FILE}' updated on Hugging Face Hub.")
    except Exception as e:
        print(f"Error updating history file on Hub: {e}")
    finally:
        if os.path.exists(temp_history_filepath):
            os.remove(temp_history_filepath)

    progress(1.0, desc="Done!")

    # ギャラリー表示用のアイテムリストを生成（Hub上のURLを使用）
    gallery_items = [(item[2], item[1].replace("|-|", "\n")) for item in history]
    
    processed_img, processed_gallery_items = process_image(img, gallery_items)
    
    latest_caption_table = make_html_table(caption_text)
    return processed_img, processed_gallery_items, latest_caption_table

import gc
import torch
def process_image(img, gallery_items): # Assuming this is part of a function
    try:
        gc.collect()        
        # Clear PyTorch's cache if GPU memory is being used
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        return img, gallery_items

    except RuntimeError as e:
        # Catch errors like CUDA Out of Memory
        error_message = f"error in generate: {e}\n\n"
        if "CUDA out of memory" in str(e):
            error_message += "memory error"
        else:
            error_message += "other error"
        print(error_message) # Output to server logs
        return None, None

# CSS 設定（ダークモード強制防止＋カフェ風テーマ）
css = """
@import url('https://fonts.googleapis.com/css2?family=Playpen+Sans+Hebrew:wght@100;200;300;400;500;600;700;800&display=swap');
body {
  background-color: #f4e1c1 !important;
  font-family:'Playpen Sans Hebrew', ‘Georgia’, serif !important;
  color: #000 !important;
}
html, .gradio-container, .dark, .dark * {
  background: #fffaf1 !important;
  color: #000 !important;
}
#col-container {
  background: #fffaf1;
  padding: 20px;
  border-radius: 16px;
  box-shadow: 0 4px 12px rgba(0,0,0,0.1);
  margin: auto;
  max-width: 780px;
}
.gr-button {
  background-color: #d4a373 !important;
  color: white !important;
  border-radius: 8px !important;
  padding: 10px 24px !important;
  font-weight: bold;
  transition: background-color 0.3s;
}
.gr-button:hover {
  background-color: #c48f61 !important;
}
.gr-textbox, .gr-slider, .gr-radio, .gr-checkbox, .gr-image {
  background: #fff;
  border-radius: 8px;
}
.gr-gallery {
  background: #fffaf1;
  padding: 10px;
  border-radius: 12px;
}
.gr-gallery .gallery-item Figcaption,
.gr-gallery .gallery-item figcaption {
  width:420px !important;
  word-wrap:break-word !important;
}
.gradio-spinner { display: none !important; }
#custom-loader {
  align-items: center;
  justify-content: center;
  font-weight: bold;
  margin: 12px 0;
  position: fixed;
  z-index: 9999;
}
.block.svelte-11xb1hd {
  background: #efd1bf !important;
}
span.svelte-g2oxp3, label.svelte-5ncdh7.svelte-5ncdh7.svelte-5ncdh7 {
  color: #915325 !important;
}
.svelte-zyxd38 g {
  display: none !important;
}
.secondary.svelte-1ixn6qd {
  background: #dca08a !important;
  color: #631c00 !important;
}
:root {
  --color-accent: #a57659;
}
.max_value.svelte-10lj3xl.svelte-10lj3xl, span.min_value {
  color: #a54618 !important;
}
body.gradio-running #custom-loader { display: flex; }
#custom-loader, .loading-text {
width: auto !important;
height: auto !important;
}
@keyframes fadeLetter {
  0%,100% { opacity: 1; }
  50%     { opacity: 0.2; }
}
#custom-loader .loading-text span {
  display: inline-block;
  animation: fadeLetter 1.8s ease-in-out infinite;
  font-size:3em;
}
#custom-loader img {
  width: 64px;
  height: 64px;
  border-radius: 50%;
  margin-left: 8px;
  display: inline-block;
  animation: jump 2s infinite ease-in-out;
  vertical-align: middle;
}
.svelte-zyxd38{
 width: 100% !important;
 height: 100% !important;
}
    @keyframes jump {
      0%, 100% { transform: translateY(10px); opacity: 1;}
      50%      { transform: translateY(-10px); opacity: 1;}
    }
.nobackground, .nobackground div, .nobackground.parent.parent.parent {
  background-color: transparent !important;
}
progress::-webkit-progress-value {
  background-color: #a57659 !important;
}
progress::-moz-progress-bar {
  background-color: #a57659 !important;
}
.gradio-progress .progress-bar,
.gradio-progress-bar {
  background-color: #a57659 !important;
}

"""

with gr.Blocks(css=css, theme=gr.themes.Default(font=[gr.themes.GoogleFont("Playpen Sans Hebrew"), "Arial", "sans-serif"])) as demo:
    with gr.Column(elem_id="col-container"):
        gr.HTML('<section class="nobackground"><h2>SDXL – Re:cocoamixXL3 (coamixXL3) Demo</h2><br>The log is shared with other. (No more than 10 images will be displayed in history.)<br>Please use this model at your own risk. I am not responsible in any way for any problems with the generated images.</section>')
        gr.HTML('<section class="nobackground"><a href="https://civitai.com/models/1553716?modelVersionId=1855218" target="_blank">Link: Civitai</a></section>')

        with gr.Row():
            prompt = gr.Textbox(lines=1, placeholder="Prompt…", value="1girl, cocoart, masterpiece, anime, high quality,", label="Prompt")
            neg    = gr.Textbox(lines=1, placeholder="Negative prompt", value="low quality, worst quality, bad, bad lighting, lowres, error, miss stroke, smoke, ugly, extra digits, creepy, imprecise, blurry,", label="Negative prompt")
        with gr.Row():
            seed_sl = gr.Slider(0, MAX_SEED, step=1, value=0, label="Seed")
            rand    = gr.Checkbox(True, label="Randomize seed")
        with gr.Row():
            width  = gr.Slider(256, 512, step=32, value=512, label="Width")
            height = gr.Slider(256, 768, step=32, value=512, label="Height")
        with gr.Row():
            cfg    = gr.Slider(1.0, 30.0, step=0.5, value=6, label="CFG Scale")
            steps  = gr.Slider(1, 15, step=1, value=12,   label="Steps")
        with gr.Row():
            scheduler_type = gr.Radio(choices=["Euler Ancestral", "DPM++ 2M SDE"], value="Euler Ancestral", label="Scheduler")
            run = gr.Button("Generate")

        # カスタムローダー
        gr.HTML(
            """
<script>
window.addEventListener('load', () => {
  const observer = new MutationObserver(() => {
    const svg = document.querySelector('svg.svelte-zyxd38');
    if (!svg) return;
    // 一度挿入したら再挿入しない
    if (svg.querySelector('foreignObject#custom-loader-fo')) return;

    const SVG_NS = 'http://www.w3.org/2000/svg';
    // foreignObject を作成
    const fo = document.createElementNS(SVG_NS, 'foreignObject');
    fo.setAttribute('id', 'custom-loader-fo');
    fo.setAttribute('width', '100%');   // SVG 内の表示エリア幅に合わせて調整
    fo.setAttribute('height', '100%');  // SVG 内の表示エリア高さに合わせて調整
    fo.setAttribute('x', '-200');
    fo.setAttribute('y', '0');

    // HTML 部分を innerHTML で一発挿入
    fo.innerHTML = `
      <div id="custom-loader" xmlns="http://www.w3.org/1999/xhtml">
        <div class="loading-text">
          <span style="animation-delay:0s">i</span>
          <span style="animation-delay:0.1s">n</span>
          <span style="animation-delay:0.2s"> </span>
          <span style="animation-delay:0.3s">p</span>
          <span style="animation-delay:0.4s">r</span>
          <span style="animation-delay:0.5s">o</span>
          <span style="animation-delay:0.6s">g</span>
          <span style="animation-delay:0.7s">r</span>
          <span style="animation-delay:0.8s">e</span>
          <span style="animation-delay:0.9s">s</span>
          <span style="animation-delay:1.0s">s</span>
          <img src="https://huggingface.co/spaces/cocoat/Re.cocoamixXL3/resolve/main/icon.png" width="32" height="32" />
        </div>
      </div>
    `;

    svg.appendChild(fo);
  });
  observer.observe(document.body, { childList: true, subtree: true });
});
</script>
            """
        )
        img_out = gr.Image(
            interactive=None, 
            value=create_dummy_image(width=512, height=512, alpha=0),
            label="生成画像" 
        )
        state   = gr.State([])
        history_gallery = gr.Gallery(
            label="生成履歴",
            columns=4,
            height=280,
            show_label=False,
            interactive=None,
            type="auto",
            value=[]
        )
        # テーブル部分だけを下にまとめて生HTMLレンダー
        history_tables = gr.HTML(value="")

    run.click(
        fn=infer,
        inputs=[prompt, neg, seed_sl, rand, width, height, cfg, steps, scheduler_type],
        outputs=[img_out, history_gallery, history_tables]
    )
    history_gallery.select(
        fn=update_history_tables_on_select,
        inputs=None, # select イベントは自動的にイベントデータ (gr.SelectData) を渡す
        outputs=[history_tables]
    )

# ページロード時に history から初期表示
    demo.load( 
        fn=lambda: ( # history リストの各要素が (Hub上のファイルパス, キャプション, Space内のファイルパス)
            [ (item[2], item[1].replace("|-|", "\n")) for item in history if item[2] is not None ], 
            make_html_table(history[0][1]) if history else "" # 最初のアイテムのキャプションを表示 
        ), 
        inputs=[], 
        outputs=[history_gallery, history_tables] 
    )
    
demo.queue()
demo.launch()