# Part 1/4 - Imports and Initial Setup
import gradio as gr
import numpy as np
import librosa
from transformers import pipeline
from datetime import datetime
import os
import requests
import json
import time

# 데이터 저장을 위한 DB 클래스
class SimpleDB:
    def __init__(self, file_path="wishes.json"):
        self.file_path = file_path
        self.wishes = self._load_wishes()

    def _load_wishes(self):
        try:
            if os.path.exists(self.file_path):
                with open(self.file_path, 'r', encoding='utf-8') as f:
                    return json.load(f)
            return []
        except Exception as e:
            print(f"Error loading wishes: {e}")
            return []

    def save_wish(self, name, wish, timestamp=None):
        if timestamp is None:
            timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        wish_data = {
            "name": name,
            "wish": wish,
            "timestamp": timestamp
        }
        self.wishes.append(wish_data)
        try:
            with open(self.file_path, 'w', encoding='utf-8') as f:
                json.dump(self.wishes, f, ensure_ascii=False, indent=2)
            return True
        except Exception as e:
            print(f"Error saving wish: {e}")
            return False

# 환경변수 및 API 설정
HF_API_TOKEN = os.getenv("roots")
if not HF_API_TOKEN:
    raise ValueError("roots token not found in environment variables")

API_URL = "https://api-inference.huggingface.co/models/stabilityai/stable-diffusion-xl-base-1.0"
headers = {"Authorization": f"Bearer {HF_API_TOKEN}"}

# AI 모델 초기화
speech_recognizer = pipeline(
    "automatic-speech-recognition",
    model="kresnik/wav2vec2-large-xlsr-korean"
)
text_analyzer = pipeline(
    "sentiment-analysis",
    model="nlptown/bert-base-multilingual-uncased-sentiment"
)

# 상수 정의
IMAGE_DISPLAY_TIME = 30
WELCOME_MESSAGE = """
# 디지털 굿판에 오신 것을 환영합니다

디지털 굿판은 현대 도시 속에서 잊혀진 전통 굿의 정수를 담아낸 **디지털 의례의 공간**입니다. 
이곳에서는 사람들의 목소리와 감정을 통해 **영적 교감**을 나누고, **자연과 도시의 에너지가 연결**됩니다. 
이제, 평온함과 치유의 여정을 시작해보세요.
"""

WORLDVIEW_MESSAGE = """
## 굿판의 세계관 🌌

온천천의 물줄기는 신성한 금샘에서 시작됩니다. 금샘은 생명과 창조의 원천이며, 
천상의 생명이 지상에서 숨을 틔우는 자리입니다. 도시의 소음 속에서도 신성한 생명력을 느껴보세요. 
이곳에서 영적인 교감을 경험하며, 자연과 하나 되는 순간을 맞이해 보시기 바랍니다.

이 앱은 온천천의 사운드스케이프를 녹음하여 제작되었으며, 
온천천 온천장역에서 장전역까지 걸으며 더 깊은 체험이 가능합니다.
"""


# Part 2/4 - Analysis Functions
def calculate_baseline_features(audio_path):
    """기준점 음성 특성 분석"""
    try:
        y, sr = librosa.load(audio_path, sr=16000)
        features = {
            "energy": float(np.mean(librosa.feature.rms(y=y))),
            "tempo": float(librosa.beat.tempo(y)[0]),
            "pitch": float(np.mean(librosa.feature.zero_crossing_rate(y))),
            "volume": float(np.mean(np.abs(y))),
            "mfcc": librosa.feature.mfcc(y=y, sr=sr, n_mfcc=13).mean(axis=1).tolist()
        }
        return features
    except Exception as e:
        print(f"Error calculating baseline: {str(e)}")
        return None

def map_acoustic_to_emotion(features, baseline_features=None):
    """음향학적 특성을 감정으로 매핑"""
    energy_norm = min(features["energy"] * 100, 100)
    tempo_norm = min(features["tempo"] / 200, 1)
    pitch_norm = min(features["pitch"] * 2, 1)

    if baseline_features:
        energy_norm = (features["energy"] / baseline_features["energy"]) * 50
        tempo_norm = (features["tempo"] / baseline_features["tempo"])
        pitch_norm = (features["pitch"] / baseline_features["pitch"])

    emotions = {
        "primary": "",
        "intensity": energy_norm,
        "confidence": 0.0,
        "secondary": "",
        "characteristics": []
    }

    if energy_norm > 70:
        if tempo_norm > 0.6:
            emotions["primary"] = "기쁨/열정"
            emotions["characteristics"].append("빠르고 활기찬 말하기 패턴")
        else:
            emotions["primary"] = "분노/강조"
            emotions["characteristics"].append("강한 음성 강도")
        emotions["confidence"] = energy_norm / 100
    elif pitch_norm > 0.6:
        if energy_norm > 50:
            emotions["primary"] = "놀람/흥분"
            emotions["characteristics"].append("높은 음고와 강한 강세")
        else:
            emotions["primary"] = "관심/호기심"
            emotions["characteristics"].append("음고 변화가 큼")
        emotions["confidence"] = pitch_norm
    elif energy_norm < 30:
        if tempo_norm < 0.4:
            emotions["primary"] = "슬픔/우울"
            emotions["characteristics"].append("느리고 약한 음성")
        else:
            emotions["primary"] = "피로/무기력"
            emotions["characteristics"].append("낮은 에너지 레벨")
        emotions["confidence"] = (30 - energy_norm) / 30
    else:
        if tempo_norm > 0.5:
            emotions["primary"] = "평온/안정"
            emotions["characteristics"].append("균형잡힌 말하기 패턴")
        else:
            emotions["primary"] = "차분/진지"
            emotions["characteristics"].append("안정적인 음성 특성")
        emotions["confidence"] = 0.5

    emotions["details"] = {
        "energy_level": f"{energy_norm:.1f}%",
        "speech_rate": f"{'빠름' if tempo_norm > 0.6 else '보통' if tempo_norm > 0.4 else '느림'}",
        "pitch_variation": f"{'높음' if pitch_norm > 0.6 else '보통' if pitch_norm > 0.3 else '낮음'}",
        "voice_volume": f"{'큼' if features['volume'] > 0.7 else '보통' if features['volume'] > 0.3 else '작음'}"
    }

    return emotions

def analyze_voice(audio_path, state):
    """통합 음성 분석"""
    if audio_path is None:
        return state, "음성을 먼저 녹음해주세요.", "", "", ""
    
    try:
        y, sr = librosa.load(audio_path, sr=16000)
        
        # 음향학적 특성 분석
        acoustic_features = {
            "energy": float(np.mean(librosa.feature.rms(y=y))),
            "tempo": float(librosa.beat.tempo(y)[0]),
            "pitch": float(np.mean(librosa.feature.zero_crossing_rate(y))),
            "volume": float(np.mean(np.abs(y)))
        }

        # 음성 감정 분석
        voice_emotion = map_acoustic_to_emotion(acoustic_features, state.get("baseline_features"))
        
        # 음성 인식
        transcription = speech_recognizer(y)
        text = transcription["text"]
        
        # 텍스트 감정 분석
        text_sentiment = text_analyzer(text)[0]
        
        # 결과 포맷팅
        voice_result = (
            f"음성 감정: {voice_emotion['primary']} "
            f"(강도: {voice_emotion['intensity']:.1f}%, 신뢰도: {voice_emotion['confidence']:.2f})\n"
            f"특징: {', '.join(voice_emotion['characteristics'])}\n"
            f"상세 분석:\n"
            f"- 에너지 레벨: {voice_emotion['details']['energy_level']}\n"
            f"- 말하기 속도: {voice_emotion['details']['speech_rate']}\n"
            f"- 음높이 변화: {voice_emotion['details']['pitch_variation']}\n"
            f"- 음성 크기: {voice_emotion['details']['voice_volume']}"
        )
        
        text_result = f"텍스트 감정 분석 (1-5): {text_sentiment['score']}"
        
        # 프롬프트 생성
        prompt = generate_detailed_prompt(text, voice_emotion, text_sentiment)
        
        return state, text, voice_result, text_result, prompt
    except Exception as e:
        return state, f"오류 발생: {str(e)}", "", "", ""


# Part 3/4 - Generation and Processing Functions
def generate_detailed_prompt(text, emotions, text_sentiment):
    """감정 기반 상세 프롬프트 생성"""
    emotion_colors = {
        "기쁨/열정": "밝은 노랑과 따뜻한 주황색",
        "분노/강조": "강렬한 빨강과 짙은 검정",
        "놀람/흥분": "선명한 파랑과 밝은 보라",
        "관심/호기심": "연한 하늘색과 민트색",
        "슬픔/우울": "어두운 파랑과 회색",
        "피로/무기력": "탁한 갈색과 짙은 회색",
        "평온/안정": "부드러운 초록과 베이지",
        "차분/진지": "차분한 남색과 깊은 보라"
    }

    if emotions["intensity"] > 70:
        visual_style = "역동적인 붓질과 강한 대비"
    elif emotions["intensity"] > 40:
        visual_style = "균형잡힌 구도와 중간 톤의 조화"
    else:
        visual_style = "부드러운 그라데이션과 차분한 톤"

    prompt = f"한국 전통 민화 스타일의 추상화, {emotion_colors.get(emotions['primary'], '자연스러운 색상')} 기반. "
    prompt += f"{visual_style}로 표현된 {emotions['primary']}의 감정. "
    prompt += f"음성의 특징({', '.join(emotions['characteristics'])})을 화면의 동적 요소로 표현. "
    prompt += f"발화 내용 '{text}'에서 느껴지는 감정(강도: {text_sentiment['score']}/5)을 은유적 이미지로 담아내기."
    
    return prompt

def generate_image_from_prompt(prompt):
    """이미지 생성 함수"""
    print(f"Generating image with prompt: {prompt}")
    try:
        if not prompt:
            print("No prompt provided")
            return None
        
        response = requests.post(
            API_URL,
            headers=headers,
            json={
                "inputs": prompt,
                "parameters": {
                    "negative_prompt": "ugly, blurry, poor quality, distorted",
                    "num_inference_steps": 30,
                    "guidance_scale": 7.5
                }
            }
        )
        
        if response.status_code == 200:
            print("Image generated successfully")
            # 이미지 저장
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            image_path = f"generated_images/{timestamp}.png"
            os.makedirs("generated_images", exist_ok=True)
            with open(image_path, "wb") as f:
                f.write(response.content)
            return response.content, image_path
        else:
            print(f"Error: {response.status_code}")
            print(f"Response: {response.text}")
            return None, None
    except Exception as e:
        print(f"Error generating image: {str(e)}")
        return None, None

def update_final_prompt(state):
    """청신의 감상들을 종합하여 최종 프롬프트 업데이트"""
    combined_prompt = "한국 전통 민화 스타일의 추상화, 온천천에서의 감상과 소원을 담아내기:\n\n"
    
    if state.get("reflections"):
        combined_prompt += "청신의 감상들:\n"
        for time, text, sentiment in state["reflections"]:
            combined_prompt += f"- {time}: {text} ({sentiment})\n"
    
    if state.get("wish"):
        combined_prompt += f"\n소원:\n{state['wish']}"
    
    return combined_prompt

def handle_image_timeout(result_image, delay=IMAGE_DISPLAY_TIME):
    """이미지 자동 사라짐 처리"""
    time.sleep(delay)
    return gr.update(value=None)

def save_reflection(text, state):
    """감상 저장"""
    if not text.strip():
        return state, state["reflections"]
    
    try:
        current_time = datetime.now().strftime("%H:%M:%S")
        sentiment = text_analyzer(text)[0]
        new_reflection = [current_time, text, f"{sentiment['label']} ({sentiment['score']:.2f})"]
        
        if "reflections" not in state:
            state["reflections"] = []
            
        state["reflections"].append(new_reflection)
        return state, state["reflections"]
    except Exception as e:
        print(f"Error in save_reflection: {str(e)}")
        return state, []


# Part 4/4 - Interface and Main
def create_interface():
    db = SimpleDB()
    
    with gr.Blocks(theme=gr.themes.Soft()) as app:
        state = gr.State({
            "user_name": "",
            "baseline_features": None,
            "reflections": [],
            "wish": None,
            "final_prompt": "",
            "image_path": None
        })

        # 헤더
        header = gr.Markdown("# 디지털 굿판")
        user_display = gr.Markdown("")

        with gr.Tabs() as tabs:
            # 입장
            with gr.Tab("입장"):
                gr.Markdown(WELCOME_MESSAGE)
                name_input = gr.Textbox(
                    label="이름을 알려주세요",
                    placeholder="이름을 입력해주세요"
                )
                worldview_display = gr.Markdown(visible=False)
                start_btn = gr.Button("여정 시작하기")

            # 기준 설정
            with gr.Tab("기준 설정"):
                gr.Markdown("""### 축원의 문장을 평온한 마음으로 읽어주세요
                
                먼저, 평온한 마음으로 축원의 문장을 읽어주세요. 
                이 축원은 당신에게 평화와 안정을 불러일으키며, 
                감정을 정확히 이해하기 위한 **기준점**이 될 것입니다.""")
                gr.Markdown("'당신의 건강과 행복이 늘 가득하기를'")
                baseline_audio = gr.Audio(
                    label="축원 문장 녹음하기",
                    sources=["microphone"]
                )
                set_baseline_btn = gr.Button("기준점 설정 완료")
                baseline_status = gr.Markdown("")

            # 청신
            with gr.Tab("청신"):
                gr.Markdown("""## 청신 - 소리로 정화하기
                
                온천천의 물소리에 귀 기울이며 **30분간 마음을 정화**해보세요.""")
                play_music_btn = gr.Button("온천천의 소리 듣기")
                with gr.Row():
                    audio = gr.Audio(
                        value=None,
                        type="filepath",
                        label="온천천의 소리",
                        interactive=False,
                        autoplay=False
                    )
                    with gr.Column():
                        reflection_input = gr.Textbox(
                            label="지금 이 순간의 감상을 자유롭게 적어보세요",
                            lines=3
                        )
                        save_btn = gr.Button("감상 저장하기")
                        reflections_display = gr.Dataframe(
                            headers=["시간", "감상", "감정 분석"],
                            label="기록된 감상들"
                        )

            # 기원
            with gr.Tab("기원"):
                gr.Markdown("## 기원 - 소원을 전해보세요")
                with gr.Row():
                    with gr.Column():
                        voice_input = gr.Audio(
                            label="소원을 나누고 싶은 마음을 말해주세요",
                            sources=["microphone"]
                        )
                        clear_btn = gr.Button("녹음 지우기")
                        analyze_btn = gr.Button("소원 분석하기")
                    
                    with gr.Column():
                        transcribed_text = gr.Textbox(label="인식된 텍스트")
                        voice_emotion = gr.Textbox(label="음성 감정 분석")
                        text_emotion = gr.Textbox(label="텍스트 감정 분석")

            # 송신
            with gr.Tab("송신"):
                gr.Markdown("## 송신 - 마음의 그림을 남기고, 보내기")
                final_prompt = gr.Textbox(
                    label="생성된 프롬프트",
                    interactive=False,
                    lines=3
                )
                with gr.Row():
                    generate_btn = gr.Button("마음의 그림 그리기")
                    save_image_btn = gr.Button("이미지 저장하기")
                result_image = gr.Image(label="생성된 이미지")
                
                gr.Markdown("## 마지막 감상을 남겨주세요")
                final_reflection = gr.Textbox(
                    label="마지막 감상",
                    placeholder="한 줄로 남겨주세요..."
                )
                save_final_btn = gr.Button("감상 남기기")

        # 이벤트 연결
        start_btn.click(
            fn=lambda name: (
                WORLDVIEW_MESSAGE if name.strip() else "이름을 입력해주세요",
                gr.update(visible=True) if name.strip() else gr.update(),
                {"user_name": name} if name.strip() else state
            ),
            inputs=[name_input],
            outputs=[worldview_display, tabs, state]
        )

        set_baseline_btn.click(
            fn=lambda x, s: ({"baseline_features": calculate_baseline_features(x)}, "기준점이 설정되었습니다."),
            inputs=[baseline_audio, state],
            outputs=[state, baseline_status]
        )

        save_btn.click(
            fn=save_reflection,
            inputs=[reflection_input, state],
            outputs=[state, reflections_display]
        )

        clear_btn.click(
            fn=lambda: None,
            outputs=[voice_input]
        )

        analyze_btn.click(
            fn=analyze_voice,
            inputs=[voice_input, state],
            outputs=[state, transcribed_text, voice_emotion, text_emotion, final_prompt]
        )

        generate_btn.click(
            fn=lambda p: generate_image_from_prompt(p)[0],
            inputs=[final_prompt],
            outputs=[result_image]
        )

        save_final_btn.click(
            fn=lambda t, s: (db.save_wish(s["user_name"], t), "감상이 저장되었습니다."),
            inputs=[final_reflection, state],
            outputs=[gr.Markdown("")]
        )

        # 이미지 자동 사라짐
        result_image.change(
            fn=lambda: gr.update(value=None),
            inputs=[],
            outputs=[result_image],
            _js=f"() => setTimeout(() => {{document.querySelector('#result_image image').src = ''}}, {IMAGE_DISPLAY_TIME*1000})"
        )

        return app

if __name__ == "__main__":
    demo = create_interface()
    demo.launch(debug=True)