"""
bridge_server.py  –  WebUI版 bridge3.py
依存: fastapi uvicorn openai

起動:
    uvicorn bridge_server:app --host 0.0.0.0 --port 8001 
"""

import os, subprocess, time, re, threading, queue, asyncio, json, zipfile, shutil
from pathlib import Path
from openai import OpenAI
from fastapi import FastAPI, HTTPException, Body
from fastapi.responses import StreamingResponse, JSONResponse, FileResponse
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel
from gradio_client import Client, handle_file
from PIL import Image
import io
import webbrowser
import base64

# --- 設定ファイル ---
SETTINGS_FILE = "bridge_settings.json"

DEFAULT_SETTINGS = {
    "tts_batch_size": 2,      # TTSに投げる文章数 (0=すべて)
    "dark_mode": False,
    "blur_mode": True,
    "tts_enabled": True,
    "bg_enabled": True,
}

def load_settings() -> dict:
    if os.path.exists(SETTINGS_FILE):
        try:
            with open(SETTINGS_FILE, "r", encoding="utf-8") as f:
                saved = json.load(f)
            merged = {**DEFAULT_SETTINGS, **saved}
            return merged
        except Exception:
            pass
    return dict(DEFAULT_SETTINGS)

def save_settings(data: dict):
    with open(SETTINGS_FILE, "w", encoding="utf-8") as f:
        json.dump(data, f, ensure_ascii=False, indent=2)

# --- 設定 ---
# LM StudioのAPIモード設定
LLM_API_URL = os.getenv("LLM_API_URL", "http://localhost:1234/v1")
LLM_API_KEY = "lm-studio"


# irodori.TTSの設定. 
TTS_API_URL = os.getenv("TTS_API_URL", "http://localhost:7862/")

# その他の設定
CHAR_DIR        = "characters"
BASE_PROMPT_FILE = "base_system_prompt.txt"
MAX_HISTORY     = 10
WINDOW_WIDTH    = 300


# TTS詳細設定
TTS_CONFIG = {
    "checkpoint": "Aratako/Irodori-TTS-500M-v2",
    "model_device": "cuda",
    "model_precision": "fp32",
    "codec_device": "cuda",
    "codec_precision": "fp32",
    "num_steps": 40,
    "num_candidates": 1,
    "seed_raw": "",
    "cfg_guidance_mode": "independent",
    "cfg_scale_text": 3,
    "cfg_scale_speaker": 5,
    "cfg_scale_raw": "",
    "cfg_min_t": 0.5,
    "cfg_max_t": 1,
    "context_kv_cache": True,
    "truncation_factor_raw": "",
    "rescale_k_raw": "",
    "rescale_sigma_raw": "",
    "speaker_kv_scale_raw": "",
    "speaker_kv_min_t_raw": "0.9",
    "speaker_kv_max_layers_raw": "",
}


# --- OpenAI クライアント ---
client = OpenAI(base_url=LLM_API_URL, api_key=LLM_API_KEY)


# 音声管理辞書
# { "req_id": {"audios": [b64,...], "queued": N, "completed": N} }
# queued == completed になったとき「全音声が届いた」と判定できる
finished_audios: dict = {}
finished_audios_lock = threading.Lock()

# --- TTS キュー ---
generate_queue: queue.Queue = queue.Queue()
error_queue = queue.Queue()


# --- 共通クライアントの保持 ---
tts_client = None

def get_tts_client():
    """Gradio Clientをシングルトンで取得"""
    global tts_client
    if tts_client is None:
        try:
            print(f"Connecting to TTS Server at {TTS_API_URL}...")
            tts_client = Client(TTS_API_URL)
            print("Successfully connected to TTS Server.")
        except Exception as e:
            print(f"Failed to connect to TTS Server: {e}")
    return tts_client


# --- TTS 生成ワーカー ---

def tts_generator_worker():
    """裏側で音声を生成し続けるワーカー"""
    while True:
        item = generate_queue.get()
        if item is None:
            break

        text, voice_path, req_id = item
        try:
            audio_b64 = generate_and_encode_tts(text, voice_path)
        except Exception as e:
            print(f"Worker Error: {e}")
            audio_b64 = None
        finally:
            # 成否に関わらず completed をインクリメントし、音声があれば追加
            with finished_audios_lock:
                if req_id in finished_audios:
                    if audio_b64:
                        finished_audios[req_id]["audios"].append(audio_b64)
                        print(f"[TTS] Generated for {req_id}: {text[:20]}...")
                    else:
                        print(f"[TTS] Failed (skipped) for {req_id}: {text[:20]}...")
                    finished_audios[req_id]["completed"] += 1
            generate_queue.task_done()

# スレッドを開始
threading.Thread(target=tts_generator_worker, daemon=True).start()


# --- キャラクター読み込み ---

IMAGE_EXTS = [".png", ".jpg", ".jpeg", ".webp"]
# zip展開キャッシュ先
ZIP_CACHE_DIR = os.path.join("characters", ".zip_cache")


def _find_image(directory: str, base: str) -> str | None:
    """directory 内で base 名に一致する画像ファイルパスを返す"""
    for ext in IMAGE_EXTS:
        p = os.path.join(directory, base + ext)
        if os.path.exists(p):
            return p
    return None


def _extract_zip(zip_path: str) -> str:
    """
    zip を ZIP_CACHE_DIR/<stem>/ に展開し、展開先ディレクトリを返す。
    既に展開済みかつ zip より新しければ再展開しない。
    """
    os.makedirs(ZIP_CACHE_DIR, exist_ok=True)
    stem = Path(zip_path).stem
    dest = os.path.join(ZIP_CACHE_DIR, stem)

    zip_mtime  = os.path.getmtime(zip_path)
    dest_mtime = os.path.getmtime(dest) if os.path.exists(dest) else 0

    if dest_mtime < zip_mtime:
        # 再展開
        if os.path.exists(dest):
            shutil.rmtree(dest)
        os.makedirs(dest, exist_ok=True)
        try:
            with zipfile.ZipFile(zip_path, "r") as zf:
                zf.extractall(dest)
            print(f"[ZIP] Extracted: {zip_path} → {dest}")
        except Exception as e:
            print(f"[ZIP] Failed to extract {zip_path}: {e}")
            shutil.rmtree(dest, ignore_errors=True)
            return ""
    return dest


def _char_from_dir(base: str, directory: str) -> dict | None:
    """
    directory 内に <base>.txt があればキャラクター辞書を返す。
    txt が存在しない場合は None。
    """
    txt_path = os.path.join(directory, base + ".txt")
    if not os.path.exists(txt_path):
        # サブフォルダ内でフォルダ名と異なるtxtを探す
        txts = [f for f in os.listdir(directory) if f.endswith(".txt")]
        if not txts:
            return None
        txt_path = os.path.join(directory, txts[0])
        base = Path(txts[0]).stem

    wav_path = os.path.join(directory, base + ".wav")
    has_voice = os.path.exists(wav_path)
    img_path  = _find_image(directory, base)

    # wav と同名でなくとも .wav が1つだけあれば使う
    if not has_voice:
        wavs = [f for f in os.listdir(directory) if f.endswith(".wav")]
        if wavs:
            wav_path  = os.path.join(directory, wavs[0])
            has_voice = True

    return {
        "name":       base if has_voice else f"{base} (音声なし)",
        "prompt":     txt_path,
        "voice_path": wav_path if has_voice else None,
        "image_path": img_path,
        "has_voice":  has_voice,
    }


def get_dynamic_characters() -> list:
    if not os.path.exists(CHAR_DIR):
        os.makedirs(CHAR_DIR)
        return []

    char_list: list[dict] = []
    seen_names: set[str]  = set()   # 重複防止

    entries = os.listdir(CHAR_DIR)

    for entry in sorted(entries):
        full_path = os.path.join(CHAR_DIR, entry)

        # ── パターン1: フラット (.txt) ──────────────────────────
        if entry.endswith(".txt"):
            base = Path(entry).stem
            if base in seen_names:
                continue
            char = _char_from_dir(base, CHAR_DIR)
            if char:
                seen_names.add(base)
                char_list.append(char)

        # ── パターン2: サブフォルダ ──────────────────────────────
        elif os.path.isdir(full_path) and entry != ".zip_cache":
            base = entry
            if base in seen_names:
                continue
            char = _char_from_dir(base, full_path)
            if char:
                seen_names.add(base)
                char_list.append(char)

        # ── パターン3: zip ───────────────────────────────────────
        elif entry.endswith(".zip"):
            base = Path(entry).stem
            if base in seen_names:
                continue
            dest = _extract_zip(full_path)
            if not dest:
                continue
            # 展開先がフォルダ1つだけの場合はその中を見る
            sub_entries = os.listdir(dest)
            if len(sub_entries) == 1 and os.path.isdir(os.path.join(dest, sub_entries[0])):
                search_dir = os.path.join(dest, sub_entries[0])
            else:
                search_dir = dest
            char = _char_from_dir(base, search_dir)
            if not char:
                # stem名と一致しなくてもtxtを探す
                char = _char_from_dir(
                    Path([f for f in os.listdir(search_dir) if f.endswith(".txt")][0]).stem
                    if any(f.endswith(".txt") for f in os.listdir(search_dir)) else base,
                    search_dir,
                )
            if char:
                seen_names.add(base)
                char_list.append(char)

    return char_list


# --- アプリ状態 ---
class AppState:
    def __init__(self):
        self.chars:          list  = []
        self.selected:       dict  = {}
        self.memory:         list  = []
        self.base_prompt:    str   = ""

state = AppState()

# --- FastAPI ---
app = FastAPI(title="Bridge WebUI")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

# 自動起動用
def open_browser():
    port = os.getenv("BRIDGE_PORT", "8001")
    url = f"http://localhost:{port}"
    time.sleep(2)
    print(f"[INFO] Opening browser at {url}")
    webbrowser.open(url)

threading.Thread(target=open_browser, daemon=True).start()


# --- index.html 配信 ---
@app.get("/")
def root():
    return FileResponse("index.html")

# --- キャラクター一覧 ---
@app.get("/characters")
def list_characters():
    state.chars = get_dynamic_characters()
    return [{
        "index": i, 
        "name": c["name"], 
        "has_image": c["image_path"] is not None,
        "has_voice": c["has_voice"]
    } for i, c in enumerate(state.chars)]

# --- キャラクター画像 ---
@app.get("/characters/{idx}/image")
def get_character_image(idx: int):
    if idx < 0 or idx >= len(state.chars):
        raise HTTPException(404)

    img_path = state.chars[idx].get("image_path")
    if not img_path or not os.path.exists(img_path):
        raise HTTPException(404)

    img = Image.open(img_path).convert("RGB")
    img = crop_center(img)
    img = img.resize((640, 640), Image.LANCZOS)

    buf = io.BytesIO()
    img.save(buf, format="JPEG", quality=95)
    buf.seek(0)

    return StreamingResponse(buf, media_type="image/jpeg")

def crop_center(img):
    w, h = img.size
    size = min(w, h)
    left = (w - size) // 2
    top = 0
    return img.crop((left, top, left + size, top + size))

# --- キャラクター選択 ---
class SelectRequest(BaseModel):
    index: int

@app.post("/select")
def select_character(req: SelectRequest):
    state.chars = get_dynamic_characters()
    if req.index < 0 or req.index >= len(state.chars):
        raise HTTPException(400, "Invalid index")

    state.selected = state.chars[req.index]

    base_prompt = ""
    if os.path.exists(BASE_PROMPT_FILE):
        with open(BASE_PROMPT_FILE, "r", encoding="utf-8") as f:
            base_prompt = f.read().strip()

    with open(state.selected["prompt"], "r", encoding="utf-8") as f:
        char_setting = f.read().strip()

    state.memory = [{"role": "system", "content": f"{base_prompt}\n\n{char_setting}"}]
    return {"status": "ok", "name": state.selected["name"]}


@app.get("/status")
async def get_status():
    is_processing = not generate_queue.empty()
    return {"is_generating": is_processing}


# --- チャット ---
class ChatRequest(BaseModel):
    message: str
    tts_enabled: bool = True


# --- 補助関数 ---
def generate_and_encode_tts(text, voice_path):
    """TTSを生成してBase64文字列で返す"""
    tts = get_tts_client()
    if tts is None:
        return None
    
    try:
        result = tts.predict(
            **TTS_CONFIG,
            text=text,
            uploaded_audio=handle_file(voice_path),
            api_name="/_run_generation"
        )
        
        res_data = result[0]
        gen_path = res_data["value"] if isinstance(res_data, dict) and "value" in res_data else res_data
        
        if gen_path and os.path.exists(gen_path):
            with open(gen_path, "rb") as f:
                return base64.b64encode(f.read()).decode('utf-8')
    except Exception as e:
        print(f"TTS Error: {e}")
    return None


def enqueue_tts(sentence: str, voice_path: str, req_id: str):
    """キューに積む前に queued カウンターをインクリメントする（必ずペアで使う）"""
    with finished_audios_lock:
        if req_id in finished_audios:
            finished_audios[req_id]["queued"] += 1
    generate_queue.put((sentence, voice_path, req_id))


EMOJI_PAT = re.compile(
    r'[\U0001F300-\U0001F9FF\U00002600-\U000027BF\U0001FA00-\U0001FFFF\U00002300-\U000023FF]+'
)

def clean_for_tts(text: str) -> str:
    """括弧内除去・絵文字除去・前後空白除去"""
    s = re.sub(r'[\(（].*?[\)）]', '', text)
    s = EMOJI_PAT.sub('', s)
    return s.strip()


def stream_llm_with_tts(messages: list, tts_enabled: bool):
    req_id = f"req_{int(time.time() * 1000)}"

    with finished_audios_lock:
        finished_audios[req_id] = {"audios": [], "queued": 0, "completed": 0}

    settings = load_settings()
    tts_batch_size: int = settings.get("tts_batch_size", 2)  # 0=すべて

    full_answer = ""
    llm_buffer = ""
    # 句読点・感嘆符の後で分割。直後の絵文字は前トークンに含めるため後続絵文字をバッファへ残す
    split_pat = re.compile(r'(?<=[。！？\?!])(?!' + EMOJI_PAT.pattern + r')')

    # バッチ用: 確定した文のリスト
    pending_sentences: list[str] = []
    enqueued_count: int = 0  # バッチ制御カウンタ

    def flush_audios():
        with finished_audios_lock:
            audios = finished_audios[req_id]["audios"]
            while audios:
                yield f"data: {json.dumps({'audio': audios.pop(0)})}\n\n"

    def maybe_enqueue_batch():
        """pending_sentences を tts_batch_size 件ずつ結合してキューに積む"""
        nonlocal pending_sentences, enqueued_count
        batch = tts_batch_size if tts_batch_size > 0 else len(pending_sentences)
        while len(pending_sentences) >= max(batch, 1):
            chunk_sentences = pending_sentences[:batch]
            pending_sentences = pending_sentences[batch:]
            combined = "".join(chunk_sentences)
            s = clean_for_tts(combined)
            if len(s) > 1:
                enqueue_tts(s, state.selected["voice_path"], req_id)
                enqueued_count += 1

    try:
        stream = client.chat.completions.create(
            model="local-model",
            messages=messages,
            stream=True,
        )

        for chunk in stream:
            delta = chunk.choices[0].delta.content or ""
            full_answer += delta
            llm_buffer += delta

            yield f"data: {json.dumps({'token': delta})}\n\n"

            if tts_enabled and state.selected.get("has_voice"):
                parts = split_pat.split(llm_buffer)
                if len(parts) > 1:
                    # 絵文字が続く場合は次のパーツの先頭絵文字を前のパーツに結合
                    merged: list[str] = []
                    for i, part in enumerate(parts[:-1]):
                        suffix_match = EMOJI_PAT.match(parts[i + 1]) if i + 1 < len(parts) else None
                        if suffix_match:
                            part = part + suffix_match.group()
                            # 次パーツから絵文字prefix を削除
                            parts[i + 1] = parts[i + 1][suffix_match.end():]
                        merged.append(part)
                    for sentence in merged:
                        if sentence.strip():
                            pending_sentences.append(sentence)
                    llm_buffer = parts[-1]
                    maybe_enqueue_batch()

            # LLM生成中に出来上がった音声はSSEでそのまま流す
            yield from flush_audios()

        # LLM終了後の残りバッファをキューへ（tts_batch_size=0 なら全部まとめて）
        if tts_enabled and state.selected.get("has_voice"):
            if llm_buffer.strip():
                pending_sentences.append(llm_buffer)
            if pending_sentences:
                # 残り全部を1バッチとして送る
                combined = "".join(pending_sentences)
                pending_sentences = []
                s = clean_for_tts(combined)
                if len(s) > 1:
                    enqueue_tts(s, state.selected["voice_path"], req_id)

        state.memory.append({"role": "assistant", "content": full_answer})

        # LLM完了を即通知。残音声があれば req_id をフロントに渡してポーリングさせる
        with finished_audios_lock:
            q = finished_audios[req_id]["queued"]
        has_pending = tts_enabled and state.selected.get("has_voice") and q > 0
        yield f"data: {json.dumps({'done': True, 'req_id': req_id if has_pending else None})}\n\n"

        # req_id を返した場合はポーリング側が解放するので、ここでは pop しない
        if not has_pending:
            with finished_audios_lock:
                finished_audios.pop(req_id, None)

    except Exception as e:
        yield f"data: {json.dumps({'error': str(e)})}\n\n"
        with finished_audios_lock:
            finished_audios.pop(req_id, None)


@app.post("/chat")
def chat(req: ChatRequest):
    if not state.selected:
        raise HTTPException(400, "キャラクターが選択されていません")

    user_text = req.message if req.message.strip() else " "
    state.memory.append({"role": "user", "content": user_text})

    return StreamingResponse(
        stream_llm_with_tts(state.memory, req.tts_enabled),
        media_type="text/event-stream"
    )


@app.get("/initial_greeting")
def initial_greeting(tts_enabled: bool = True):
    if not state.selected:
        raise HTTPException(400, "Character not selected")
        
    prompt = "会話開始時の最初の一言として、短く自然な挨拶をしてください。"
    state.memory.append({"role": "user", "content": prompt})

    return StreamingResponse(
        stream_llm_with_tts(state.memory, tts_enabled),
        media_type="text/event-stream"
    )


@app.get("/audio_poll/{req_id}")
def audio_poll(req_id: str):
    """
    バックグラウンドTTS生成の進捗をポーリングするエンドポイント。
    - audios: 今この瞬間に出来上がっている音声(base64)リスト
    - finished: queued==completed で全音声生成完了
    """
    with finished_audios_lock:
        entry = finished_audios.get(req_id)
        if entry is None:
            return JSONResponse({"audios": [], "finished": True})

        audios = entry["audios"][:]
        entry["audios"].clear()
        q = entry["queued"]
        c = entry["completed"]
        finished = (q > 0 and q == c)

    if finished:
        # 全完了したらメモリ解放
        with finished_audios_lock:
            finished_audios.pop(req_id, None)
        print(f"[TTS] All {q} audio(s) polled and released for {req_id}")

    return JSONResponse({"audios": audios, "finished": finished})


@app.on_event("startup")
def on_startup():
    threading.Thread(target=open_browser, daemon=True).start()

@app.get("/settings")
def get_settings():
    """設定ファイルを読み込んで返す"""
    return JSONResponse(load_settings())


@app.post("/settings")
def post_settings(data: dict = Body(...)):
    """設定を保存する"""
    current = load_settings()
    current.update(data)
    save_settings(current)
    return JSONResponse({"status": "ok"})
