#!/usr/bin/env python3

import os
import re
import tempfile
import gc
from collections.abc import Iterator
from threading import Thread, Lock
import json
import requests
import cv2
import gradio as gr
import spaces
import torch
import numpy as np
from loguru import logger
from PIL import Image
from transformers import AutoProcessor, Gemma3ForConditionalGeneration, TextIteratorStreamer, pipeline
import time
import warnings
from typing import Dict, List, Optional, Union
import librosa
import scipy.signal as sps
import queue

# CSV/TXT 분석
import pandas as pd
# PDF 텍스트 추출
import PyPDF2

warnings.filterwarnings('ignore')

# 로깅 설정
logger.remove()
logger.add(lambda msg: print(msg, flush=True), level="INFO")

print("🎮 로봇 시각 시스템 초기화 (Gemma3-R1984-4B + Whisper)...")

##############################################################################
# 상수 정의
##############################################################################
MAX_CONTENT_CHARS = 2000
MAX_INPUT_LENGTH = 2096
MAX_NUM_IMAGES = 5
SERPHOUSE_API_KEY = os.getenv("SERPHOUSE_API_KEY", "")

##############################################################################
# 전역 변수
##############################################################################
model = None
processor = None
whisper_model = None
model_loaded = False
whisper_loaded = False
model_name = "Gemma3-R1984-4B"

# 오디오 관련 전역 변수
audio_lock = Lock()
last_audio_data = None
last_transcription = ""

##############################################################################
# 메모리 관리
##############################################################################
def clear_cuda_cache():
    """CUDA 캐시를 명시적으로 비웁니다."""
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
        gc.collect()

##############################################################################
# 키워드 추출 함수
##############################################################################
def extract_keywords(text: str, top_k: int = 5) -> str:
    """키워드 추출"""
    text = re.sub(r"[^a-zA-Z0-9가-힣\s]", "", text)
    tokens = text.split()
    
    seen = set()
    unique_tokens = []
    for token in tokens:
        if token not in seen and len(token) > 1:
            seen.add(token)
            unique_tokens.append(token)
    
    key_tokens = unique_tokens[:top_k]
    return " ".join(key_tokens)

##############################################################################
# Whisper 모델 로드
##############################################################################
@spaces.GPU(duration=60)
def load_whisper():
    global whisper_model, whisper_loaded
    
    if whisper_loaded:
        logger.info("Whisper 모델이 이미 로드되어 있습니다.")
        return True
    
    try:
        logger.info("Whisper 모델 로딩 시작...")
        # 파이프라인 방식으로 로드
        device = 0 if torch.cuda.is_available() else "cpu"
        whisper_model = pipeline(
            task="automatic-speech-recognition",
            model="openai/whisper-base",
            chunk_length_s=30,
            device=device,
        )
        whisper_loaded = True
        logger.info("✅ Whisper 모델 로딩 완료!")
        return True
    except Exception as e:
        logger.error(f"Whisper 모델 로딩 실패: {e}")
        return False

##############################################################################
# 오디오 처리 함수 (간소화)
##############################################################################
def resample_audio(audio: np.ndarray, orig_sr: int, target_sr: int = 16000) -> np.ndarray:
    """오디오 리샘플링"""
    if orig_sr == target_sr:
        return audio.astype(np.float32)
    
    # scipy를 사용한 리샘플링
    number_of_samples = round(len(audio) * float(target_sr) / orig_sr)
    audio_resampled = sps.resample(audio, number_of_samples)
    return audio_resampled.astype(np.float32)

@spaces.GPU(duration=30)
def transcribe_audio_whisper(audio_array: np.ndarray, sr: int = 16000):
    """Whisper를 사용한 오디오 전사"""
    global whisper_model, whisper_loaded
    
    if not whisper_loaded:
        if not load_whisper():
            return None
    
    try:
        # 오디오가 너무 조용한지 체크
        if np.max(np.abs(audio_array)) < 0.01:
            logger.warning("오디오가 너무 조용함")
            return None
            
        # 음성 인식
        result = whisper_model({"array": audio_array, "sampling_rate": sr})
        transcription = result["text"].strip()
        
        logger.info(f"Whisper 전사 성공: {transcription[:50]}...")
        return transcription if transcription else None
        
    except Exception as e:
        logger.error(f"Whisper 오디오 전사 오류: {e}")
        import traceback
        logger.error(traceback.format_exc())
        return None

def process_audio_recording(audio_data):
    """녹음된 오디오 처리"""
    global last_audio_data, last_transcription, audio_lock
    
    if audio_data is None:
        return None
    
    try:
        # 오디오 데이터 추출
        if isinstance(audio_data, tuple) and len(audio_data) == 2:
            sr, audio = audio_data
        else:
            logger.warning(f"예상치 못한 오디오 형식: {type(audio_data)}")
            return None
        
        if audio is None or len(audio) == 0:
            return None
        
        # numpy 배열로 변환
        if not isinstance(audio, np.ndarray):
            audio = np.array(audio)
        
        # 스테레오를 모노로 변환
        if audio.ndim > 1:
            audio = audio.mean(axis=1)
        
        # 16kHz로 리샘플링
        if sr != 16000:
            audio = resample_audio(audio, sr, 16000)
        
        # 저장
        with audio_lock:
            last_audio_data = (audio, 16000)
        
        logger.info(f"오디오 저장 완료: {len(audio)/16000:.1f}초")
        
        # 전사 시도
        transcription = transcribe_audio_whisper(audio, 16000)
        if transcription:
            with audio_lock:
                last_transcription = transcription
            return transcription
        
    except Exception as e:
        logger.error(f"오디오 처리 오류: {e}")
        import traceback
        logger.error(traceback.format_exc())
    
    return None

##############################################################################
# 웹 검색 함수
##############################################################################
def do_web_search(query: str) -> str:
    """SerpHouse API를 사용한 웹 검색"""
    try:
        url = "https://api.serphouse.com/serp/live"
        
        params = {
            "q": query,
            "domain": "google.com",
            "serp_type": "web",
            "device": "desktop",
            "lang": "ko",  # 한국어 우선
            "num": "10"   # 10개로 제한
        }
        
        headers = {
            "Authorization": f"Bearer {SERPHOUSE_API_KEY}"
        }
        
        logger.info(f"웹 검색 중... 검색어: {query}")
        
        response = requests.get(url, headers=headers, params=params, timeout=60)
        response.raise_for_status()
        
        data = response.json()
        
        results = data.get("results", {})
        organic = results.get("organic", []) if isinstance(results, dict) else []
        
        if not organic:
            return "검색 결과를 찾을 수 없습니다."
        
        max_results = min(10, len(organic))
        limited_organic = organic[:max_results]
        
        summary_lines = []
        for idx, item in enumerate(limited_organic, start=1):
            title = item.get("title", "제목 없음")
            link = item.get("link", "#")
            snippet = item.get("snippet", "설명 없음")
            displayed_link = item.get("displayed_link", link)
            
            summary_lines.append(
                f"### 결과 {idx}: {title}\n\n"
                f"{snippet}\n\n"
                f"**출처**: [{displayed_link}]({link})\n\n"
                f"---\n"
            )
        
        instructions = """# 웹 검색 결과
아래는 검색 결과입니다. 답변 시 이 정보를 활용하세요:
1. 각 결과의 제목, 내용, 출처 링크를 참조하세요
2. 관련 출처를 명시적으로 인용하세요
3. 여러 출처의 정보를 종합하여 답변하세요
"""
        
        search_results = instructions + "\n".join(summary_lines)
        return search_results
    
    except Exception as e:
        logger.error(f"웹 검색 실패: {e}")
        return f"웹 검색 실패: {str(e)}"

##############################################################################
# 문서 처리 함수
##############################################################################
def analyze_csv_file(path: str) -> str:
    """CSV 파일 분석"""
    try:
        df = pd.read_csv(path)
        if df.shape[0] > 50 or df.shape[1] > 10:
            df = df.iloc[:50, :10]
        df_str = df.to_string()
        if len(df_str) > MAX_CONTENT_CHARS:
            df_str = df_str[:MAX_CONTENT_CHARS] + "\n...(중략)..."
        return f"**[CSV 파일: {os.path.basename(path)}]**\n\n{df_str}"
    except Exception as e:
        return f"CSV 읽기 실패 ({os.path.basename(path)}): {str(e)}"

def analyze_txt_file(path: str) -> str:
    """TXT 파일 분석"""
    try:
        with open(path, "r", encoding="utf-8") as f:
            text = f.read()
        if len(text) > MAX_CONTENT_CHARS:
            text = text[:MAX_CONTENT_CHARS] + "\n...(중략)..."
        return f"**[TXT 파일: {os.path.basename(path)}]**\n\n{text}"
    except Exception as e:
        return f"TXT 읽기 실패 ({os.path.basename(path)}): {str(e)}"

def pdf_to_markdown(pdf_path: str) -> str:
    """PDF를 마크다운으로 변환"""
    text_chunks = []
    try:
        with open(pdf_path, "rb") as f:
            reader = PyPDF2.PdfReader(f)
            max_pages = min(5, len(reader.pages))
            for page_num in range(max_pages):
                page = reader.pages[page_num]
                page_text = page.extract_text() or ""
                page_text = page_text.strip()
                if page_text:
                    if len(page_text) > MAX_CONTENT_CHARS // max_pages:
                        page_text = page_text[:MAX_CONTENT_CHARS // max_pages] + "...(중략)"
                    text_chunks.append(f"## 페이지 {page_num+1}\n\n{page_text}\n")
            if len(reader.pages) > max_pages:
                text_chunks.append(f"\n...({max_pages}/{len(reader.pages)} 페이지 표시)...")
    except Exception as e:
        return f"PDF 읽기 실패 ({os.path.basename(pdf_path)}): {str(e)}"

    full_text = "\n".join(text_chunks)
    if len(full_text) > MAX_CONTENT_CHARS:
        full_text = full_text[:MAX_CONTENT_CHARS] + "\n...(중략)..."

    return f"**[PDF 파일: {os.path.basename(pdf_path)}]**\n\n{full_text}"

##############################################################################
# 모델 로드
##############################################################################
@spaces.GPU(duration=120)
def load_model():
    global model, processor, model_loaded
    
    if model_loaded:
        logger.info("모델이 이미 로드되어 있습니다.")
        return True
    
    try:
        logger.info("Gemma3-R1984-4B 모델 로딩 시작...")
        clear_cuda_cache()
        
        model_id = os.getenv("MODEL_ID", "VIDraft/Gemma-3-R1984-4B")
        
        processor = AutoProcessor.from_pretrained(model_id, padding_side="left")
        model = Gemma3ForConditionalGeneration.from_pretrained(
            model_id,
            device_map="auto",
            torch_dtype=torch.bfloat16,
            attn_implementation="eager"
        )
        
        model_loaded = True
        logger.info(f"✅ {model_name} 로딩 완료!")
        return True
        
    except Exception as e:
        logger.error(f"모델 로딩 실패: {e}")
        return False

##############################################################################
# 이미지 분석 (로봇 태스크 중심)
##############################################################################
@spaces.GPU(duration=60)
def analyze_image_for_robot(
    image: Union[np.ndarray, Image.Image],
    prompt: str,
    task_type: str = "general",
    use_web_search: bool = False,
    enable_thinking: bool = False,
    max_new_tokens: int = 300,
    audio_transcript: Optional[str] = None
) -> str:
    """로봇 작업을 위한 이미지 분석 (오디오 정보 포함)"""
    global model, processor
    
    if not model_loaded:
        if not load_model():
            return "❌ 모델 로딩 실패"
    
    try:
        # numpy 배열을 PIL 이미지로 변환
        if isinstance(image, np.ndarray):
            image = Image.fromarray(image).convert('RGB')
        
        # 태스크별 시스템 프롬프트 구성
        system_prompts = {
            "general": "당신은 로봇 시각 시스템입니다. 먼저 장면을 1-2줄로 설명하고, 핵심 내용을 간결하게 분석하세요.",
            "planning": """당신은 로봇 작업 계획 AI입니다. 
먼저 장면 이해를 1-2줄로 설명하고, 그 다음 작업 계획을 작성하세요.
형식: 
[장면 이해] 현재 보이는 장면을 1-2줄로 설명

[작업 계획]
Step_1: xxx
Step_2: xxx
Step_n: xxx""",
            "grounding": "당신은 객체 위치 시스템입니다. 먼저 보이는 객체들을 한 줄로 설명하고, 요청된 객체 위치를 [x1, y1, x2, y2]로 반환하세요.",
            "affordance": "당신은 파지점 분석 AI입니다. 먼저 대상 객체를 한 줄로 설명하고, 파지 영역을 [x1, y1, x2, y2]로 반환하세요.",
            "trajectory": "당신은 경로 계획 AI입니다. 먼저 환경을 한 줄로 설명하고, 경로를 [(x1,y1), (x2,y2), ...]로 제시하세요.",
            "pointing": "당신은 지점 지정 시스템입니다. 먼저 참조점들을 한 줄로 설명하고, 위치를 [(x1,y1), (x2,y2), ...]로 반환하세요."
        }
        
        # 오디오 정보가 있으면 프롬프트 수정
        if audio_transcript and task_type == "planning":
            system_prompts["planning"] = """당신은 로봇 작업 계획 AI입니다. 
먼저 장면 이해를 1-2줄로 설명하고, 주변 소리를 인식했다면 그것도 설명한 후, 작업 계획을 작성하세요.
형식: 
[장면 이해] 현재 보이는 장면을 1-2줄로 설명
[주변 소리 인식] 들리는 소리나 음성을 1줄로 설명
[작업 계획]
Step_1: xxx
Step_2: xxx
Step_n: xxx"""
        
        system_prompt = system_prompts.get(task_type, system_prompts["general"])
        
        # Chain-of-Thought 추가 (선택적)
        if enable_thinking:
            system_prompt += "\n\n추론 과정을 <thinking></thinking> 태그 안에 작성 후 최종 답변을 제시하세요. 장면 이해는 추론 과정과 별도로 반드시 포함하세요."
        
        # 웹 검색 수행
        combined_system = system_prompt
        if use_web_search:
            keywords = extract_keywords(prompt, top_k=5)
            if keywords:
                logger.info(f"웹 검색 키워드: {keywords}")
                search_results = do_web_search(keywords)
                combined_system = f"{search_results}\n\n{system_prompt}"
        
        # 사용자 프롬프트에 오디오 정보 추가
        user_prompt = prompt
        if audio_transcript:
            user_prompt += f"\n\n[인식된 주변 소리: {audio_transcript}]"
        
        # 메시지 구성
        messages = [
            {
                "role": "system",
                "content": [{"type": "text", "text": combined_system}]
            },
            {
                "role": "user",
                "content": [
                    {"type": "image", "url": image},
                    {"type": "text", "text": user_prompt}
                ]
            }
        ]
        
        # 입력 처리
        inputs = processor.apply_chat_template(
            messages,
            add_generation_prompt=True,
            tokenize=True,
            return_dict=True,
            return_tensors="pt",
        ).to(device=model.device, dtype=torch.bfloat16)
        
        # 입력 토큰 수 제한
        if inputs.input_ids.shape[1] > MAX_INPUT_LENGTH:
            inputs.input_ids = inputs.input_ids[:, -MAX_INPUT_LENGTH:]
            if 'attention_mask' in inputs:
                inputs.attention_mask = inputs.attention_mask[:, -MAX_INPUT_LENGTH:]
        
        # 생성
        with torch.no_grad():
            outputs = model.generate(
                **inputs,
                max_new_tokens=max_new_tokens,
                do_sample=True,
                temperature=0.7,
                top_p=0.9,
                pad_token_id=processor.tokenizer.pad_token_id,
                eos_token_id=processor.tokenizer.eos_token_id,
            )
        
        # 입력 토큰 제거하여 출력만 추출
        generated_tokens = outputs[0][inputs.input_ids.shape[1]:]
        
        # 디코딩
        response = processor.decode(generated_tokens, skip_special_tokens=True).strip()
        
        # 프롬프트 제거 및 정리
        # 이미 입력 토큰을 제거했으므로 추가 정리만 수행
        response = response.strip()
        
        # 혹시 남아있는 불필요한 텍스트 제거
        if response.startswith("model\n"):
            response = response[6:].strip()
        elif response.startswith("model"):
            response = response[5:].strip()
        
        return response
        
    except Exception as e:
        logger.error(f"이미지 분석 오류: {e}")
        import traceback
        return f"❌ 분석 오류: {str(e)}\n{traceback.format_exc()}"
    finally:
        clear_cuda_cache()

##############################################################################
# 문서 분석 (스트리밍)
##############################################################################
def _model_gen_with_oom_catch(**kwargs):
    """OOM 처리를 위한 생성 함수"""
    global model
    try:
        model.generate(**kwargs)
    except torch.cuda.OutOfMemoryError:
        raise RuntimeError("GPU 메모리 부족. Max Tokens를 줄여주세요.")
    finally:
        clear_cuda_cache()

@spaces.GPU(duration=120)
def analyze_documents_streaming(
    files: List[str],
    prompt: str,
    use_web_search: bool = False,
    max_new_tokens: int = 2048
) -> Iterator[str]:
    """문서 분석 (스트리밍)"""
    global model, processor
    
    if not model_loaded:
        if not load_model():
            yield "❌ 모델 로딩 실패"
            return
    
    try:
        # 시스템 프롬프트
        system_content = "당신은 문서를 분석하고 요약하는 전문 AI입니다."
        
        # 웹 검색
        if use_web_search:
            keywords = extract_keywords(prompt, top_k=5)
            if keywords:
                search_results = do_web_search(keywords)
                system_content = f"{search_results}\n\n{system_content}"
        
        # 문서 내용 처리
        doc_contents = []
        for file_path in files:
            if file_path.lower().endswith('.csv'):
                content = analyze_csv_file(file_path)
            elif file_path.lower().endswith('.txt'):
                content = analyze_txt_file(file_path)
            elif file_path.lower().endswith('.pdf'):
                content = pdf_to_markdown(file_path)
            else:
                continue
            doc_contents.append(content)
        
        # 메시지 구성
        messages = [
            {
                "role": "system",
                "content": [{"type": "text", "text": system_content}]
            },
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": "\n\n".join(doc_contents) + f"\n\n{prompt}"}
                ]
            }
        ]
        
        # 입력 처리
        inputs = processor.apply_chat_template(
            messages,
            add_generation_prompt=True,
            tokenize=True,
            return_dict=True,
            return_tensors="pt",
        ).to(device=model.device, dtype=torch.bfloat16)
        
        # 스트리밍 설정
        streamer = TextIteratorStreamer(processor, timeout=30.0, skip_prompt=True, skip_special_tokens=True)
        gen_kwargs = dict(
            inputs,
            streamer=streamer,
            max_new_tokens=max_new_tokens,
            temperature=0.8,
            top_p=0.9,
        )
        
        # 별도 스레드에서 생성
        t = Thread(target=_model_gen_with_oom_catch, kwargs=gen_kwargs)
        t.start()
        
        # 스트리밍 출력
        output = ""
        for new_text in streamer:
            output += new_text
            yield output
            
    except Exception as e:
        logger.error(f"문서 분석 오류: {e}")
        yield f"❌ 오류 발생: {str(e)}"
    finally:
        clear_cuda_cache()

##############################################################################
# Gradio UI (로봇 시각화 중심)
##############################################################################
css = """
.robot-header {
    text-align: center;
    background: linear-gradient(135deg, #1e3c72 0%, #2a5298 50%, #667eea 100%);
    color: white;
    padding: 20px;
    border-radius: 10px;
    margin-bottom: 20px;
    box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1);
}
.status-box {
    text-align: center;
    padding: 10px;
    border-radius: 5px;
    margin: 10px 0;
    font-weight: bold;
}
.info-box {
    background: #f0f0f0;
    padding: 15px;
    border-radius: 8px;
    margin: 10px 0;
    border-left: 4px solid #2a5298;
}
.task-button {
    min-height: 60px;
    font-size: 1.1em;
}
.webcam-container {
    border: 3px solid #2a5298;
    border-radius: 10px;
    padding: 10px;
    background: #f8f9fa;
}
.auto-capture-status {
    text-align: center;
    padding: 5px;
    border-radius: 5px;
    margin: 5px 0;
    font-weight: bold;
    background: #e8f5e9;
    color: #2e7d32;
}
.audio-status {
    text-align: center;
    padding: 5px;
    border-radius: 5px;
    margin: 5px 0;
    font-weight: bold;
    background: #e3f2fd;
    color: #1565c0;
}
"""

with gr.Blocks(title="🤖 로봇 시각 시스템 (Gemma3-4B)", css=css) as demo:
    gr.HTML("""
    <div class="robot-header">
        <h1>🤖 로봇 시각 시스템</h1>
        <h3>🎮 Gemma3-R1984-4B + 📷 실시간 웹캠 + 🎤 음성 인식</h3>
        <p>⚡ 멀티모달 AI로 로봇 작업 분석!</p>
    </div>
    """)
    
    with gr.Row():
        # 왼쪽: 웹캠 및 입력
        with gr.Column(scale=1):
            gr.Markdown("### 📷 실시간 웹캠")
            
            with gr.Group(elem_classes="webcam-container"):
                webcam = gr.Image(
                    sources=["webcam"], 
                    streaming=True,
                    type="numpy",
                    label="실시간 스트리밍",
                    height=300
                )
                
                # 자동 캡처 상태 표시
                auto_capture_status = gr.HTML(
                    '<div class="auto-capture-status">🔄 자동 캡처: 대기 중</div>'
                )
                
                # 캡처된 이미지 표시
                captured_image = gr.Image(
                    label="캡처된 이미지",
                    height=180,
                    visible=False
                )
            
            # 오디오 컨트롤
            gr.Markdown("### 🎤 음성 인식")
            with gr.Group():
                # 오디오 상태 표시
                audio_status = gr.HTML(
                    '<div class="audio-status">🎤 음성 인식: 비활성화</div>'
                )
                
                # 녹음 인터페이스 (숨김 상태로 시작)
                audio_recorder = gr.Audio(
                    sources=["microphone"],
                    type="numpy",
                    label="🎤 10초 녹음",
                    visible=False
                )
                
                # 마지막 인식된 텍스트
                last_transcript = gr.Textbox(
                    label="인식된 음성",
                    value="",
                    lines=2,
                    interactive=False
                )
            
            # 로봇 작업 버튼들
            gr.Markdown("### 🎯 로봇 작업")
            with gr.Row():
                capture_btn = gr.Button("📸 수동 캡처", variant="primary", elem_classes="task-button")
                clear_capture_btn = gr.Button("🗑️ 초기화", elem_classes="task-button")
            
            with gr.Column():
                auto_capture_toggle = gr.Checkbox(
                    label="🔄 자동 캡처 (10초마다)",
                    value=False
                )
                
                use_audio_toggle = gr.Checkbox(
                    label="🎤 음성 인식 사용",
                    value=False,
                    info="10초마다 음성을 인식하여 분석에 포함"
                )
            
            with gr.Row():
                planning_btn = gr.Button("📋 작업 계획", elem_classes="task-button")
                grounding_btn = gr.Button("📍 객체 위치", elem_classes="task-button")
        
        # 오른쪽: 분석 설정 및 결과
        with gr.Column(scale=2):
            gr.Markdown("### ⚙️ 분석 설정")
            
            with gr.Row():
                with gr.Column():
                    task_prompt = gr.Textbox(
                        label="작업 설명",
                        placeholder="예: 테이블 위의 컵을 잡아서 싱크대에 놓기",
                        value="현재 장면을 분석하고 로봇이 수행할 수 있는 작업을 제안하세요.",
                        lines=2
                    )
                    
                    with gr.Row():
                        use_web_search = gr.Checkbox(
                            label="🔍 웹 검색",
                            value=False
                        )
                        
                        enable_thinking = gr.Checkbox(
                            label="🤔 추론 과정",
                            value=False
                        )
                    
                    max_tokens = gr.Slider(
                        label="최대 토큰",
                        minimum=100,
                        maximum=1000,
                        value=300,
                        step=50
                    )
            
            gr.Markdown("### 📊 분석 결과")
            result_output = gr.Textbox(
                label="AI 분석 결과",
                lines=18,
                max_lines=35,
                show_copy_button=True,
                elem_id="result"
            )
            
            status_display = gr.HTML(
                '<div class="status-box" style="background:#d4edda; color:#155724;">🎮 시스템 준비</div>'
            )
    
    # 문서 분석 탭 (숨김)
    with gr.Tab("📄 문서 분석", visible=False):
        with gr.Row():
            with gr.Column():
                doc_files = gr.File(
                    label="문서 업로드",
                    file_count="multiple",
                    file_types=[".pdf", ".csv", ".txt"],
                    type="filepath"
                )
                
                doc_prompt = gr.Textbox(
                    label="분석 요청",
                    placeholder="예: 이 문서들의 핵심 내용을 요약하고 비교 분석하세요.",
                    lines=3
                )
                
                doc_web_search = gr.Checkbox(
                    label="🔍 웹 검색 사용",
                    value=False
                )
                
                analyze_docs_btn = gr.Button("📊 문서 분석", variant="primary")
            
            with gr.Column():
                doc_result = gr.Textbox(
                    label="분석 결과",
                    lines=25,
                    max_lines=50
                )
    
    # 이벤트 핸들러
    webcam_state = gr.State(None)
    
    def capture_webcam(frame):
        """웹캠 프레임 캡처"""
        if frame is None:
            return None, None, '<div class="status-box" style="background:#f8d7da; color:#721c24;">❌ 웹캠 프레임 없음</div>'
        return frame, gr.update(value=frame, visible=True), '<div class="status-box" style="background:#d4edda; color:#155724;">✅ 이미지 캡처 완료</div>'
    
    def clear_capture():
        """캡처 초기화"""
        global last_transcription, last_audio_data, audio_lock
        
        with audio_lock:
            last_transcription = ""
            last_audio_data = None
            
        return None, gr.update(visible=False), '<div class="status-box" style="background:#d4edda; color:#155724;">🎮 시스템 준비</div>', ""
    
    def analyze_with_task(image, prompt, task_type, use_search, thinking, tokens):
        """특정 태스크로 이미지 분석"""
        global last_transcription, audio_lock
        
        if image is None:
            return "❌ 먼저 이미지를 캡처하세요.", '<div class="status-box" style="background:#f8d7da; color:#721c24;">❌ 이미지 없음</div>'
        
        status = f'<div class="status-box" style="background:#cce5ff; color:#004085;">🚀 {task_type} 분석 중...</div>'
        
        # 현재 전사 텍스트 가져오기
        transcript = ""
        with audio_lock:
            transcript = last_transcription
        
        result = analyze_image_for_robot(
            image=image,
            prompt=prompt,
            task_type=task_type,
            use_web_search=use_search,
            enable_thinking=thinking,
            max_new_tokens=tokens,
            audio_transcript=transcript if transcript else None
        )
        
        # 결과 포맷팅
        timestamp = time.strftime("%H:%M:%S")
        task_names = {
            "planning": "작업 계획",
            "grounding": "객체 위치",
            "affordance": "파지점",
            "trajectory": "경로 계획"
        }
        
        formatted_result = f"""🤖 {task_names.get(task_type, '분석')} 결과 ({timestamp})
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
{result}
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"""
        
        complete_status = '<div class="status-box" style="background:#d4edda; color:#155724;">✅ 분석 완료!</div>'
        return formatted_result, complete_status
    
    # 자동 캡처 및 분석 함수
    @spaces.GPU(duration=60)
    def auto_capture_and_analyze(webcam_frame, task_prompt, use_search, thinking, tokens, use_audio, audio_data):
        """자동 캡처 및 분석"""
        global last_transcription, audio_lock
        
        if webcam_frame is None:
            return (
                None,
                "자동 캡처 대기 중...",
                '<div class="status-box" style="background:#fff3cd; color:#856404;">⏳ 웹캠 대기 중</div>',
                '<div class="auto-capture-status">🔄 자동 캡처: 웹캠 대기 중</div>',
                "대기 중...",
                None  # 오디오 리셋
            )
        
        # 캡처 수행
        timestamp = time.strftime("%H:%M:%S")
        
        # 오디오 처리 (있으면)
        if use_audio and audio_data is not None:
            logger.info(f"[{timestamp}] 오디오 처리 시작")
            transcription = process_audio_recording(audio_data)
            if transcription:
                logger.info(f"새로운 전사: {transcription[:50]}...")
        
        # 마지막 전사 결과 가져오기
        audio_transcript = ""
        if use_audio:
            with audio_lock:
                audio_transcript = last_transcription
                if audio_transcript:
                    logger.info(f"분석에 사용할 음성: {audio_transcript[:50]}...")
        
        # 이미지 분석 (작업 계획 모드로)
        result = analyze_image_for_robot(
            image=webcam_frame,
            prompt=task_prompt,
            task_type="planning",
            use_web_search=use_search,
            enable_thinking=thinking,
            max_new_tokens=tokens,
            audio_transcript=audio_transcript if audio_transcript else None
        )
        
        formatted_result = f"""🔄 자동 분석 완료 ({timestamp})
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
{result}
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"""
        
        # 마지막 인식된 텍스트 업데이트
        transcript_display = audio_transcript if audio_transcript else "음성 인식 대기 중..."
        
        return (
            webcam_frame,
            formatted_result,
            '<div class="status-box" style="background:#d4edda; color:#155724;">✅ 자동 분석 완료</div>',
            f'<div class="auto-capture-status">🔄 자동 캡처: 마지막 분석 {timestamp}</div>',
            transcript_display,
            None  # 오디오 리셋 (다음 녹음 준비)
        )
    
    # 웹캠 스트리밍
    webcam.stream(
        fn=lambda x: x,
        inputs=[webcam],
        outputs=[webcam_state]
    )
    
    # 수동 캡처 버튼
    capture_btn.click(
        fn=capture_webcam,
        inputs=[webcam_state],
        outputs=[webcam_state, captured_image, status_display]
    )
    
    # 초기화 버튼
    clear_capture_btn.click(
        fn=clear_capture,
        outputs=[webcam_state, captured_image, status_display, last_transcript]
    )
    
    # 작업 버튼들
    planning_btn.click(
        fn=lambda img, p, s, t, tk: analyze_with_task(img, p, "planning", s, t, tk),
        inputs=[captured_image, task_prompt, use_web_search, enable_thinking, max_tokens],
        outputs=[result_output, status_display]
    )
    
    grounding_btn.click(
        fn=lambda img, p, s, t, tk: analyze_with_task(img, p, "grounding", s, t, tk),
        inputs=[captured_image, task_prompt, use_web_search, enable_thinking, max_tokens],
        outputs=[result_output, status_display]
    )
    
    # 문서 분석
    def analyze_docs(files, prompt, use_search):
        if not files:
            return "❌ 문서를 업로드하세요."
        
        output = ""
        for chunk in analyze_documents_streaming(files, prompt, use_search):
            output = chunk
        return output
    
    analyze_docs_btn.click(
        fn=analyze_docs,
        inputs=[doc_files, doc_prompt, doc_web_search],
        outputs=[doc_result]
    )
    
    # 자동 캡처 타이머 (10초마다)
    timer = gr.Timer(10.0, active=False)
    
    # 자동 캡처 토글 이벤트
    def toggle_auto_capture(enabled):
        if enabled:
            return gr.Timer(10.0, active=True), '<div class="auto-capture-status">🔄 자동 캡처: 활성화됨 (10초마다)</div>'
        else:
            return gr.Timer(active=False), '<div class="auto-capture-status">🔄 자동 캡처: 비활성화됨</div>'
    
    auto_capture_toggle.change(
        fn=toggle_auto_capture,
        inputs=[auto_capture_toggle],
        outputs=[timer, auto_capture_status]
    )
    
    # 오디오 토글 이벤트
    def toggle_audio(enabled):
        global last_transcription, last_audio_data, audio_lock
        
        if enabled:
            # Whisper 모델 로드
            load_whisper()
            
            # 초기화
            with audio_lock:
                last_transcription = ""
                last_audio_data = None
            
            logger.info("오디오 인식 활성화됨")
            
            return (
                gr.update(visible=True),  # audio_recorder 표시
                '<div class="audio-status">🎤 음성 인식: 활성화됨</div>'
            )
        else:
            # 초기화
            with audio_lock:
                last_transcription = ""
                last_audio_data = None
            
            logger.info("오디오 인식 비활성화됨")
                
            return (
                gr.update(visible=False),  # audio_recorder 숨김
                '<div class="audio-status">🎤 음성 인식: 비활성화</div>'
            )
    
    use_audio_toggle.change(
        fn=toggle_audio,
        inputs=[use_audio_toggle],
        outputs=[audio_recorder, audio_status]
    )
    
    # 오디오 녹음 완료 시 처리
    def on_audio_recorded(audio_data):
        """오디오 녹음 완료 시 자동 처리"""
        global last_transcription, audio_lock
        
        if audio_data is not None:
            logger.info("새 오디오 녹음 감지")
            transcription = process_audio_recording(audio_data)
            if transcription:
                return transcription
        
        with audio_lock:
            return last_transcription if last_transcription else "음성 인식 대기 중..."
    
    audio_recorder.change(
        fn=on_audio_recorded,
        inputs=[audio_recorder],
        outputs=[last_transcript]
    )
    
    # 타이머 틱 이벤트
    timer.tick(
        fn=auto_capture_and_analyze,
        inputs=[webcam_state, task_prompt, use_web_search, enable_thinking, max_tokens, use_audio_toggle, audio_recorder],
        outputs=[captured_image, result_output, status_display, auto_capture_status, last_transcript, audio_recorder]
    )
    
    # 초기 모델 로드
    def initial_load():
        load_model()
        return "시스템 준비 완료! 🚀"
    
    demo.load(
        fn=initial_load,
        outputs=None
    )

if __name__ == "__main__":
    print("🚀 로봇 시각 시스템 시작 (Gemma3-R1984-4B + Whisper)...")
    demo.queue().launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
        show_error=True,
        debug=False
    )