import os
import gradio as gr
from openai import OpenAI
from PIL import Image
from io import BytesIO
import traceback
import requests

#os.environ["OPENAI_API_KEY"] = "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"
#os.environ["STABILITY_API_KEY"] = "sk-HvTCuoGI1JjUoJD2GoJ4xi7BDycodYLEIFFU6t3nJDORusLb"
# ===== 1. 환경 변수에서 API 키 읽기 =====
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
STABILITY_API_KEY = os.getenv("STABILITY_API_KEY")

if OPENAI_API_KEY is None:
    raise ValueError("OPENAI_API_KEY 환경 변수가 설정되지 않았습니다.")
if STABILITY_API_KEY is None:
    raise ValueError("STABILITY_API_KEY 환경 변수가 설정되지 않았습니다.")

client = OpenAI(api_key=OPENAI_API_KEY)
MODEL = "gpt-4o-mini"

STABILITY_ENDPOINT = "https://api.stability.ai/v2beta/stable-image/generate/core"

# ===== 2. 프롬프트 기법별 System Prompt =====
SYS_COT = """당신은 시상식 경험이 많은 전문 스타일리스트 AI입니다.
- 내부 사고과정은 숨기고 결과 요약과 핵심 근거만 출력.
- 예산/시간/동선/조명/체형/피부톤/드레스코드 반영, 한국어로 답변.
"""

SYS_TOT = """당신은 '패션 전략가'입니다.
- 서로 다른 3가지 스타일 경로 제시(각 5~7줄) → 간단 평가표(1~5) → 최종안.
- 내부 나뭇가지 추론은 숨기고 결과만 출력, 한국어.
"""

SYS_SELFCONS = """당신은 '스타일 제안 컨설턴트'입니다.
- 관점이 다른 3안 제시(클래식/모던/대담) → 기준별(무대 적합성/사진발/착용 난이도/리스크) 점수화 → 최종안 선택.
- 내부 합의 과정은 숨기고 결과만 출력, 한국어.
"""

SYS_REACT = """당신은 '스타일링 오퍼레이션 매니저'입니다.
- Plan → Action → Observation → Update 2~3회 후 Final 제안.
- 내부 사고는 숨기고 각 단계 1~3줄, 한국어.
"""

METHOD_TO_SYS = {
    "Chain-of-Thought(요약형)": SYS_COT,
    "Tree-of-Thought": SYS_TOT,
    "Self-Consistency": SYS_SELFCONS,
    "ReAct": SYS_REACT,
}

# ===== 3. RCCF 프롬프트 빌더 & 텍스트 스타일 추천 =====
def build_rccf_prompt(role, context, constraints, fmt):
    return f"""[Role]
{role}
[Context]
{context}
[Constraints]
{constraints}
[Format]
{fmt}
"""

def run_chat(method_name, role, context, constraints, fmt, model=MODEL, temperature=0.7):
    if method_name not in METHOD_TO_SYS:
        raise ValueError("기법 선택이 올바르지 않습니다.")
    system_prompt = METHOD_TO_SYS[method_name]
    user_prompt = build_rccf_prompt(role, context, constraints, fmt)

    resp = client.chat.completions.create(
        model=model,
        messages=[
            {"role": "system", "content": system_prompt},
            {"role": "user",   "content": user_prompt},
        ],
        temperature=float(temperature),
    )
    return f"### 기법: **{method_name}**\n\n" + resp.choices[0].message.content

# ===== 4. 한국어 스타일 텍스트 → 영어 이미지 프롬프트 =====
def build_image_prompt_from_text(style_text: str) -> str:
    system_msg = (
        "You are a fashion illustration prompt engineer for an image model. "
        "Given a Korean styling recommendation, output ONE concise English prompt "
        "for a full-body fashion illustration for a red carpet or award show. "
        "40 words or less. Focus on outfit, colors, silhouette, and overall vibe. "
        "No explanations, just the prompt."
    )

    resp = client.chat.completions.create(
        model=MODEL,
        messages=[
            {"role": "system", "content": system_msg},
            {"role": "user", "content": style_text},
        ],
        temperature=0.4,
    )
    prompt_en = resp.choices[0].message.content.strip()
    return prompt_en

# ===== 5. Stability AI로 이미지 생성 =====
def generate_fashion_image_with_stability(style_text: str) -> Image.Image:
    img_prompt = build_image_prompt_from_text(style_text)

    headers = {
        "Authorization": f"Bearer {STABILITY_API_KEY}",
        "Accept": "image/*",
    }

    files = {
        "prompt": (None, img_prompt),
        "output_format": (None, "png"),
    }

    response = requests.post(
        STABILITY_ENDPOINT,
        headers=headers,
        files=files,
        timeout=120,
    )

    if response.status_code != 200:
        raise RuntimeError(
            f"Stability API 오류: {response.status_code} - {response.text[:500]}"
        )

    img = Image.open(BytesIO(response.content))
    return img

# ===== 6. Gradio 인터페이스 함수 =====
def interface_fn(role, context, constraints, fmt, method, temperature):
    try:
        style_text = run_chat(method, role, context, constraints, fmt, MODEL, temperature)
        style_image = generate_fashion_image_with_stability(style_text)
        return style_text, style_image
    except Exception as e:
        traceback.print_exc()
        err_text = (
            f"⚠ 오류가 발생했습니다: {e}\n\n"
            f"Space Logs에서 traceback을 확인해 주세요."
        )
        return err_text, None

# ===== 7. Gradio UI 정의 =====
demo = gr.Interface(
    fn=interface_fn,
    inputs=[
        gr.Textbox(label="Role", value="시상식 전문 스타일리스트", lines=2),
        gr.Textbox(
            label="Context",
            value="배우 한서연이 청룡영화상 레드카펫 참석. 뉴트럴-웜 톤, 키 167cm. 미니멀 & 모던 선호.",
            lines=3,
        ),
        gr.Textbox(
            label="Constraints",
            value="이동 동선 촉박, 2벌 피팅, 헤어/메이크업 각 50분, 예산 1,000만원 협찬 우선",
            lines=3,
        ),
        gr.Textbox(
            label="Format",
            value=(
                "1) 룩 콘셉트(한 줄)\n"
                "2) 의상 제안(드레스 2안 + 대안 수트 1안)\n"
                "3) 액세서리·슈즈·클러치\n"
                "4) 헤어 & 메이크업(키포인트 3개)\n"
                "5) 리스크 & 완화(최대 3개)\n"
                "6) 현장 타임라인 체크리스트(6단계)"
            ),
            lines=8,
        ),
        gr.Dropdown(
            label="프롬프트 기법",
            choices=list(METHOD_TO_SYS.keys()),
            value="Chain-of-Thought(요약형)",
        ),
        gr.Slider(0.0, 1.2, value=0.7, step=0.1, label="Temperature"),
    ],
    outputs=[
        gr.Markdown(label="스타일 추천 결과"),
        gr.Image(label="패션 일러스트", type="pil"),
    ],
    title="나만의 스타일리스트 챗봇 — 텍스트(OpenAI) + 이미지(Stability AI)",
    description="RCCF + 프롬프트 기법으로 스타일을 추천하고, 그 결과를 바탕으로 Stable Diffusion으로 패션 일러스트 이미지를 생성합니다.",
)

# HF Spaces에서는 launch()에 share 필요 X
if __name__ == "__main__":
    demo.launch()