#!/usr/bin/env python3

import os
import re
import tempfile
import gc
from collections.abc import Iterator
from threading import Thread
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
import time
import warnings
from typing import Dict, List, Optional, Union
import base64
from io import BytesIO

# llama-cpp-python for GGUF
from llama_cpp import Llama
from llama_cpp.llama_chat_format import Llava16ChatHandler

# Model download
from huggingface_hub import hf_hub_download

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

warnings.filterwarnings('ignore')

print("🎮 로봇 시각 시스템 초기화 (Gemma3-4B GGUF Q4_K_M)...")

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

##############################################################################
# 전역 변수
##############################################################################
llm = None
model_loaded = False
model_name = "Gemma3-4B-GGUF-Q4_K_M"

##############################################################################
# 메모리 관리
##############################################################################
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)

##############################################################################
# 웹 검색 함수
##############################################################################
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"
        }
        
        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}"

##############################################################################
# 이미지를 base64로 변환
##############################################################################
def image_to_base64_data_uri(image: Union[np.ndarray, Image.Image]) -> str:
    """이미지를 base64 data URI로 변환"""
    if isinstance(image, np.ndarray):
        image = Image.fromarray(image).convert('RGB')
    
    buffered = BytesIO()
    image.save(buffered, format="JPEG", quality=85)
    img_str = base64.b64encode(buffered.getvalue()).decode()
    return f"data:image/jpeg;base64,{img_str}"

##############################################################################
# 모델 로드
##############################################################################
def download_model_files():
    """Hugging Face Hub에서 모델 파일 다운로드"""
    # 여러 가능한 저장소 시도
    model_repos = [
        # 첫 번째 시도: 일반적인 Gemma 3 4B GGUF
        {
            "repo": "Mungert/gemma-3-4b-it-gguf",
            "model": "google_gemma-3-4b-it-q4_k_m.gguf",
            "mmproj": "google_gemma-3-4b-it-mmproj-bf16.gguf"
        },
        # 두 번째 시도: LM Studio 버전
        {
            "repo": "lmstudio-community/gemma-3-4b-it-GGUF",
            "model": "gemma-3-4b-it-Q4_K_M.gguf",
            "mmproj": "gemma-3-4b-it-mmproj-f16.gguf"
        },
        # 세 번째 시도: unsloth 버전
        {
            "repo": "unsloth/gemma-3-4b-it-GGUF",
            "model": "gemma-3-4b-it.Q4_K_M.gguf",
            "mmproj": "gemma-3-4b-it.mmproj.gguf"
        }
    ]
    
    for repo_info in model_repos:
        try:
            logger.info(f"저장소 시도: {repo_info['repo']}")
            
            # 메인 모델 다운로드
            model_filename = repo_info["model"]
            logger.info(f"모델 다운로드 중: {model_filename}")
            
            model_path = hf_hub_download(
                repo_id=repo_info["repo"],
                filename=model_filename,
                resume_download=True,
                local_files_only=False
            )
            
            # Vision projection 파일 다운로드
            mmproj_filename = repo_info["mmproj"]
            logger.info(f"Vision 모델 다운로드 중: {mmproj_filename}")
            
            try:
                mmproj_path = hf_hub_download(
                    repo_id=repo_info["repo"],
                    filename=mmproj_filename,
                    resume_download=True,
                    local_files_only=False
                )
            except:
                # mmproj 파일이 없을 수도 있음
                logger.warning(f"Vision 모델을 찾을 수 없습니다: {mmproj_filename}")
                logger.warning("텍스트 전용 모드로 진행합니다.")
                mmproj_path = None
            
            logger.info(f"✅ 모델 다운로드 성공!")
            logger.info(f"모델 경로: {model_path}")
            if mmproj_path:
                logger.info(f"Vision 경로: {mmproj_path}")
            
            return model_path, mmproj_path
            
        except Exception as e:
            logger.error(f"저장소 {repo_info['repo']} 시도 실패: {e}")
            continue
    
    # 모든 시도가 실패한 경우
    raise Exception("사용 가능한 GGUF 모델을 찾을 수 없습니다. 인터넷 연결을 확인하세요.")

@spaces.GPU(duration=120)
def load_model():
    global llm, model_loaded
    
    if model_loaded:
        logger.info("모델이 이미 로드되어 있습니다.")
        return True
    
    try:
        logger.info("Gemma3-4B GGUF Q4_K_M 모델 로딩 시작...")
        clear_cuda_cache()
        
        # 모델 파일 다운로드
        model_path, mmproj_path = download_model_files()
        
        # GPU 사용 가능 여부 확인
        n_gpu_layers = -1 if torch.cuda.is_available() else 0
        
        # 채팅 핸들러 생성 (비전 지원 - mmproj가 있는 경우만)
        chat_handler = None
        if mmproj_path:
            try:
                chat_handler = Llava16ChatHandler(
                    clip_model_path=mmproj_path,
                    verbose=False
                )
                logger.info("✅ Vision 모델 로드 성공")
            except Exception as e:
                logger.warning(f"Vision 모델 로드 실패, 텍스트 전용 모드로 전환: {e}")
                chat_handler = None
        
        # 모델 로드
        llm_params = {
            "model_path": model_path,
            "n_ctx": 4096,  # 컨텍스트 크기
            "n_gpu_layers": n_gpu_layers,  # GPU 레이어
            "n_threads": 8,  # CPU 스레드
            "verbose": False,
            "seed": 42,
        }
        
        # chat_handler가 있으면 추가
        if chat_handler:
            llm_params["chat_handler"] = chat_handler
            llm_params["logits_all"] = True  # 비전 모델에 필요
        
        llm = Llama(**llm_params)
        
        model_loaded = True
        logger.info(f"✅ Gemma3-4B 모델 로딩 완료!")
        if not chat_handler:
            logger.warning("⚠️ 텍스트 전용 모드로 실행 중입니다. 이미지 분석이 제한될 수 있습니다.")
        return True
        
    except Exception as e:
        logger.error(f"모델 로딩 실패: {e}")
        import traceback
        logger.error(traceback.format_exc())
        return False

##############################################################################
# 채팅 템플릿 포맷팅
##############################################################################
def format_chat_prompt(system_prompt: str, user_prompt: str, image_uri: Optional[str] = None) -> List[Dict]:
    """Gemma 스타일 채팅 프롬프트 생성"""
    messages = []
    
    # 시스템 메시지
    messages.append({
        "role": "system",
        "content": system_prompt
    })
    
    # 사용자 메시지
    user_content = []
    if image_uri:
        user_content.append({
            "type": "image_url",
            "image_url": {"url": image_uri}
        })
    user_content.append({
        "type": "text",
        "text": user_prompt
    })
    
    messages.append({
        "role": "user",
        "content": user_content
    })
    
    return messages

##############################################################################
# 이미지 분석 (로봇 태스크 중심)
##############################################################################
@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
) -> str:
    """로봇 작업을 위한 이미지 분석"""
    global llm
    
    if not model_loaded:
        if not load_model():
            return "❌ 모델 로딩 실패"
    
    try:
        # Vision 모델이 없는 경우 경고
        if not hasattr(llm, 'chat_handler') or llm.chat_handler is None:
            logger.warning("Vision 모델이 로드되지 않았습니다. 텍스트 기반 분석만 가능합니다.")
            
            # 텍스트 전용 분석
            system_prompt = f"""당신은 로봇 시각 시스템 시뮬레이터입니다. 
실제 이미지를 볼 수는 없지만, 사용자의 설명을 바탕으로 로봇 작업을 계획하고 분석합니다.
태스크 유형: {task_type}"""
            
            messages = [
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": f"[이미지 분석 요청] {prompt}"}
            ]
            
            response = llm.create_chat_completion(
                messages=messages,
                max_tokens=max_new_tokens,
                temperature=0.7,
                top_p=0.9,
                stream=False
            )
            
            result = response['choices'][0]['message']['content'].strip()
            return f"⚠️ 텍스트 전용 모드\n\n{result}"
        
        # 이미지를 base64로 변환
        image_uri = image_to_base64_data_uri(image)
        
        # 태스크별 시스템 프롬프트 구성
        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), ...]로 반환하세요."
        }
        
        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}"
        
        # 메시지 구성
        messages = format_chat_prompt(combined_system, prompt, image_uri)
        
        # 생성
        response = llm.create_chat_completion(
            messages=messages,
            max_tokens=max_new_tokens,
            temperature=0.7,
            top_p=0.9,
            stream=False
        )
        
        # 응답 추출
        result = response['choices'][0]['message']['content'].strip()
        
        return result
        
    except Exception as e:
        logger.error(f"이미지 분석 오류: {e}")
        import traceback
        return f"❌ 분석 오류: {str(e)}\n{traceback.format_exc()}"
    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 llm
    
    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)
        
        # 전체 프롬프트 구성
        full_prompt = "\n\n".join(doc_contents) + f"\n\n{prompt}"
        
        # 메시지 구성
        messages = [
            {"role": "system", "content": system_content},
            {"role": "user", "content": full_prompt}
        ]
        
        # 스트리밍 생성
        stream = llm.create_chat_completion(
            messages=messages,
            max_tokens=max_new_tokens,
            temperature=0.8,
            top_p=0.9,
            stream=True
        )
        
        # 스트리밍 출력
        output = ""
        for chunk in stream:
            if 'choices' in chunk and len(chunk['choices']) > 0:
                delta = chunk['choices'][0].get('delta', {})
                if 'content' in delta:
                    output += delta['content']
                    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;
}
.model-info {
    background: #fff3cd;
    color: #856404;
    padding: 10px;
    border-radius: 5px;
    margin: 10px 0;
    text-align: center;
}
"""

with gr.Blocks(title="🤖 로봇 시각 시스템 (Gemma3-4B GGUF)", css=css) as demo:
    gr.HTML("""
    <div class="robot-header">
        <h1>🤖 로봇 시각 시스템</h1>
        <h3>🎮 Gemma3-4B GGUF Q4_K_M + 📷 실시간 웹캠 + 🔍 웹 검색</h3>
        <p>⚡ 양자화 모델로 더 빠르고 효율적인 로봇 작업 분석!</p>
    </div>
    """)
    
    gr.HTML("""
    <div class="model-info">
        <strong>모델:</strong> Gemma3-4B Q4_K_M (2.5GB) | <strong>메모리 사용:</strong> ~3-4GB VRAM
    </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=350
                )
                
                # 자동 캡처 상태 표시
                auto_capture_status = gr.HTML(
                    '<div class="auto-capture-status">🔄 자동 캡처: 대기 중</div>'
                )
                
                # 캡처된 이미지 표시
                captured_image = gr.Image(
                    label="캡처된 이미지",
                    height=200,
                    visible=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.Row():
                auto_capture_toggle = gr.Checkbox(
                    label="🔄 자동 캡처 활성화 (10초마다)",
                    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.Row():
                affordance_btn = gr.Button("🤏 파지점 분석", elem_classes="task-button")
                trajectory_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,
                            info="관련 정보를 웹에서 검색합니다"
                        )
                        
                        enable_thinking = gr.Checkbox(
                            label="🤔 추론 과정 표시",
                            value=False,
                            info="Chain-of-Thought 추론 과정을 보여줍니다"
                        )
                    
                    max_tokens = gr.Slider(
                        label="최대 토큰 수",
                        minimum=100,
                        maximum=2048,
                        value=300,
                        step=50
                    )
            
            gr.Markdown("### 📊 분석 결과")
            result_output = gr.Textbox(
                label="AI 분석 결과",
                lines=20,
                max_lines=40,
                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)
    auto_capture_state = gr.State({"enabled": False, "timer": 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():
        """캡처 초기화"""
        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):
        """특정 태스크로 이미지 분석"""
        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>'
        
        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
        )
        
        # 결과 포맷팅
        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
    
    # 자동 캡처 및 분석 함수
    def auto_capture_and_analyze(webcam_frame, task_prompt, use_search, thinking, tokens, auto_state):
        """자동 캡처 및 분석"""
        if webcam_frame is None:
            return (
                None,
                "자동 캡처 대기 중...",
                '<div class="status-box" style="background:#fff3cd; color:#856404;">⏳ 웹캠 대기 중</div>',
                '<div class="auto-capture-status">🔄 자동 캡처: 웹캠 대기 중</div>'
            )
        
        # 캡처 수행
        timestamp = time.strftime("%H:%M:%S")
        
        # 이미지 분석 (작업 계획 모드로)
        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
        )
        
        formatted_result = f"""🔄 자동 분석 완료 ({timestamp})
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
{result}
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"""
        
        return (
            webcam_frame,
            formatted_result,
            '<div class="status-box" style="background:#d4edda; color:#155724;">✅ 자동 분석 완료</div>',
            f'<div class="auto-capture-status">🔄 자동 캡처: 마지막 분석 {timestamp}</div>'
        )
    
    # 웹캠 스트리밍
    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]
    )
    
    # 작업 버튼들
    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]
    )
    
    affordance_btn.click(
        fn=lambda img, p, s, t, tk: analyze_with_task(img, p, "affordance", s, t, tk),
        inputs=[captured_image, task_prompt, use_web_search, enable_thinking, max_tokens],
        outputs=[result_output, status_display]
    )
    
    trajectory_btn.click(
        fn=lambda img, p, s, t, tk: analyze_with_task(img, p, "trajectory", 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]
    )
    
    # 타이머 틱 이벤트
    timer.tick(
        fn=auto_capture_and_analyze,
        inputs=[webcam_state, task_prompt, use_web_search, enable_thinking, max_tokens, auto_capture_state],
        outputs=[captured_image, result_output, status_display, auto_capture_status]
    )
    
    # 초기 모델 로드
    def initial_load():
        # 첫 실행 시 GPU에서 모델 로드
        return "시스템 준비 완료! 첫 분석 시 모델이 자동으로 로드됩니다. 🚀"
    
    demo.load(
        fn=initial_load,
        outputs=None
    )

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