import gradio as gr
import httpx
import asyncio
import json
from typing import Generator

# API 기본 URL
API_BASE_URL = "https://vidraft-kt.hf.space"

# 현재 세션 정보 저장
current_session = {"session_id": None, "user_email": None}


async def create_session(user_email: str) -> str:
    """새 세션 생성"""
    if not user_email:
        return "❌ 이메일을 입력해주세요."
    
    try:
        async with httpx.AsyncClient(timeout=30.0) as client:
            response = await client.post(
                f"{API_BASE_URL}/session/new",
                json={"user_email": user_email}
            )
            result = response.json()
            
            if "error" in result:
                return f"❌ 오류: {result['error']}"
            
            current_session["session_id"] = result.get("session_id")
            current_session["user_email"] = user_email
            return f"✅ 세션 생성 완료!\n세션 ID: {result.get('session_id')}"
    except Exception as e:
        return f"❌ 연결 오류: {str(e)}"


async def end_session(session_id: str, user_email: str) -> str:
    """세션 종료"""
    if not session_id or not user_email:
        return "❌ 세션 ID와 이메일이 필요합니다."
    
    try:
        async with httpx.AsyncClient(timeout=30.0) as client:
            response = await client.post(
                f"{API_BASE_URL}/session/end",
                json={"session_id": session_id, "user_email": user_email}
            )
            result = response.json()
            return f"✅ 세션 종료 완료: {json.dumps(result, ensure_ascii=False, indent=2)}"
    except Exception as e:
        return f"❌ 오류: {str(e)}"


def chat_streaming(message: str, user_email: str, session_id: str, 
                   web_search: bool, self_learning: bool, history: list) -> Generator:
    """스트리밍 채팅 - messages 형식 사용 (딕셔너리)"""
    if not user_email:
        history = history or []
        history.append({"role": "assistant", "content": "❌ 이메일을 먼저 입력해주세요."})
        yield history
        return
    
    if not session_id:
        history = history or []
        history.append({"role": "assistant", "content": "❌ 세션을 먼저 생성해주세요."})
        yield history
        return
    
    # 사용자 메시지 추가
    history = history or []
    history.append({"role": "user", "content": message})
    history.append({"role": "assistant", "content": ""})
    yield history
    
    # 스트리밍 응답
    full_response = ""
    try:
        with httpx.Client(timeout=120.0) as client:
            with client.stream(
                "POST",
                f"{API_BASE_URL}/chat/text/stream",
                json={
                    "message": message,
                    "user_email": user_email,
                    "session_id": session_id,
                    "web_search_enabled": web_search,
                    "self_learning_enabled": self_learning
                }
            ) as response:
                for line in response.iter_lines():
                    if line.startswith("data: "):
                        data = line[6:]
                        if data == "[DONE]":
                            break
                        try:
                            chunk = json.loads(data)
                            content = chunk.get("content", "")
                            full_response += content
                            history[-1]["content"] = full_response
                            yield history
                        except json.JSONDecodeError:
                            continue
    except Exception as e:
        history[-1]["content"] = f"❌ 오류: {str(e)}"
        yield history


async def get_health_current(user_email: str) -> str:
    """현재 건강 상태 조회"""
    if not user_email:
        return "❌ 이메일을 입력해주세요."
    
    try:
        async with httpx.AsyncClient(timeout=30.0) as client:
            response = await client.get(
                f"{API_BASE_URL}/health/current",
                params={"user_email": user_email}
            )
            result = response.json()
            return json.dumps(result, ensure_ascii=False, indent=2)
    except Exception as e:
        return f"❌ 오류: {str(e)}"


async def get_health_history(user_email: str, days: int) -> str:
    """건강 히스토리 조회"""
    if not user_email:
        return "❌ 이메일을 입력해주세요."
    
    try:
        async with httpx.AsyncClient(timeout=30.0) as client:
            response = await client.get(
                f"{API_BASE_URL}/health/history",
                params={"user_email": user_email, "days": days}
            )
            result = response.json()
            return json.dumps(result, ensure_ascii=False, indent=2)
    except Exception as e:
        return f"❌ 오류: {str(e)}"


async def search_memories(query: str, user_email: str, k: int, threshold: float) -> str:
    """기억 검색"""
    if not user_email or not query:
        return "❌ 이메일과 검색어를 입력해주세요."
    
    try:
        async with httpx.AsyncClient(timeout=30.0) as client:
            response = await client.post(
                f"{API_BASE_URL}/memory/search",
                json={
                    "query": query,
                    "user_email": user_email,
                    "k": k,
                    "threshold": threshold
                }
            )
            result = response.json()
            return json.dumps(result, ensure_ascii=False, indent=2)
    except Exception as e:
        return f"❌ 오류: {str(e)}"


async def get_all_memories(user_email: str) -> str:
    """모든 기억 조회"""
    if not user_email:
        return "❌ 이메일을 입력해주세요."
    
    try:
        async with httpx.AsyncClient(timeout=30.0) as client:
            response = await client.get(
                f"{API_BASE_URL}/memory/all",
                params={"user_email": user_email}
            )
            result = response.json()
            return json.dumps(result, ensure_ascii=False, indent=2)
    except Exception as e:
        return f"❌ 오류: {str(e)}"


async def delete_memory(memory_id: int, user_email: str) -> str:
    """기억 삭제"""
    if not user_email or not memory_id:
        return "❌ 이메일과 기억 ID를 입력해주세요."
    
    try:
        async with httpx.AsyncClient(timeout=30.0) as client:
            response = await client.delete(
                f"{API_BASE_URL}/memory/{int(memory_id)}",
                params={"user_email": user_email}
            )
            result = response.json()
            return f"✅ 삭제 완료: {json.dumps(result, ensure_ascii=False)}"
    except Exception as e:
        return f"❌ 오류: {str(e)}"


async def web_search(query: str, count: int) -> str:
    """웹 검색"""
    if not query:
        return "❌ 검색어를 입력해주세요."
    
    try:
        async with httpx.AsyncClient(timeout=30.0) as client:
            response = await client.post(
                f"{API_BASE_URL}/search/web",
                json={"query": query, "count": count}
            )
            result = response.json()
            return json.dumps(result, ensure_ascii=False, indent=2)
    except Exception as e:
        return f"❌ 오류: {str(e)}"


async def crawl_url(url: str) -> str:
    """URL 크롤링"""
    if not url:
        return "❌ URL을 입력해주세요."
    
    try:
        async with httpx.AsyncClient(timeout=60.0) as client:
            response = await client.post(
                f"{API_BASE_URL}/crawl/url",
                json={"url": url}
            )
            result = response.json()
            return json.dumps(result, ensure_ascii=False, indent=2)
    except Exception as e:
        return f"❌ 오류: {str(e)}"


async def get_storage_info() -> str:
    """스토리지 정보 조회"""
    try:
        async with httpx.AsyncClient(timeout=30.0) as client:
            response = await client.get(f"{API_BASE_URL}/storage/info")
            result = response.json()
            return json.dumps(result, ensure_ascii=False, indent=2)
    except Exception as e:
        return f"❌ 오류: {str(e)}"


# Gradio UI 구성
with gr.Blocks(title="KT API 테스트") as demo:
    gr.Markdown("# 🧪 KT 시니어 AI 비서 API 테스트")
    gr.Markdown(f"**API URL:** `{API_BASE_URL}`")
    
    with gr.Tabs():
        # 탭 1: 채팅 테스트
        with gr.TabItem("💬 채팅"):
            with gr.Row():
                with gr.Column(scale=1):
                    email_input = gr.Textbox(label="이메일 (필수)", placeholder="test@example.com")
                    session_input = gr.Textbox(label="세션 ID", placeholder="세션 생성 후 자동 입력됨")
                    
                    with gr.Row():
                        create_btn = gr.Button("🆕 세션 생성", variant="primary")
                        end_btn = gr.Button("🔚 세션 종료")
                    
                    session_result = gr.Textbox(label="세션 결과", lines=3)
                    
                    web_search_check = gr.Checkbox(label="웹 검색 활성화", value=True)
                    self_learning_check = gr.Checkbox(label="자가 학습 활성화", value=True)
                
                with gr.Column(scale=2):
                    # messages 형식 사용 (기본값)
                    chatbot = gr.Chatbot(label="대화", height=400)
                    msg_input = gr.Textbox(label="메시지 입력", placeholder="메시지를 입력하세요...")
                    send_btn = gr.Button("전송", variant="primary")
            
            # 세션 생성 이벤트
            async def on_create_session(email):
                result = await create_session(email)
                session_id = current_session.get("session_id", "")
                return result, session_id
            
            create_btn.click(
                on_create_session,
                inputs=[email_input],
                outputs=[session_result, session_input]
            )
            
            # 세션 종료 이벤트
            end_btn.click(
                end_session,
                inputs=[session_input, email_input],
                outputs=[session_result]
            )
            
            # 채팅 이벤트
            send_btn.click(
                chat_streaming,
                inputs=[msg_input, email_input, session_input, web_search_check, self_learning_check, chatbot],
                outputs=[chatbot]
            ).then(lambda: "", outputs=[msg_input])
            
            msg_input.submit(
                chat_streaming,
                inputs=[msg_input, email_input, session_input, web_search_check, self_learning_check, chatbot],
                outputs=[chatbot]
            ).then(lambda: "", outputs=[msg_input])
        
        # 탭 2: 건강 상태
        with gr.TabItem("🏥 건강 상태"):
            health_email = gr.Textbox(label="이메일", placeholder="test@example.com")
            
            with gr.Row():
                health_current_btn = gr.Button("현재 건강 상태 조회", variant="primary")
                health_days = gr.Slider(1, 90, value=30, step=1, label="히스토리 일수")
                health_history_btn = gr.Button("건강 히스토리 조회")
            
            health_result = gr.Textbox(label="결과", lines=15)
            
            health_current_btn.click(get_health_current, inputs=[health_email], outputs=[health_result])
            health_history_btn.click(get_health_history, inputs=[health_email, health_days], outputs=[health_result])
        
        # 탭 3: 기억 관리
        with gr.TabItem("🧠 기억 관리"):
            memory_email = gr.Textbox(label="이메일", placeholder="test@example.com")
            
            with gr.Row():
                with gr.Column():
                    gr.Markdown("### 기억 검색")
                    search_query = gr.Textbox(label="검색어")
                    search_k = gr.Slider(1, 50, value=10, step=1, label="결과 개수")
                    search_threshold = gr.Slider(0.1, 1.0, value=0.7, step=0.1, label="유사도 임계값")
                    search_btn = gr.Button("검색", variant="primary")
                
                with gr.Column():
                    gr.Markdown("### 기억 삭제")
                    delete_id = gr.Number(label="기억 ID", precision=0)
                    delete_btn = gr.Button("삭제", variant="stop")
                    get_all_btn = gr.Button("모든 기억 조회")
            
            memory_result = gr.Textbox(label="결과", lines=15)
            
            search_btn.click(search_memories, inputs=[search_query, memory_email, search_k, search_threshold], outputs=[memory_result])
            get_all_btn.click(get_all_memories, inputs=[memory_email], outputs=[memory_result])
            delete_btn.click(delete_memory, inputs=[delete_id, memory_email], outputs=[memory_result])
        
        # 탭 4: 웹 검색 & 크롤링
        with gr.TabItem("🌐 웹 검색"):
            with gr.Row():
                with gr.Column():
                    gr.Markdown("### 웹 검색")
                    web_query = gr.Textbox(label="검색어")
                    web_count = gr.Slider(1, 20, value=10, step=1, label="결과 개수")
                    web_search_btn = gr.Button("검색", variant="primary")
                
                with gr.Column():
                    gr.Markdown("### URL 크롤링")
                    crawl_url_input = gr.Textbox(label="URL", placeholder="https://example.com")
                    crawl_btn = gr.Button("크롤링", variant="primary")
            
            web_result = gr.Textbox(label="결과", lines=15)
            
            web_search_btn.click(web_search, inputs=[web_query, web_count], outputs=[web_result])
            crawl_btn.click(crawl_url, inputs=[crawl_url_input], outputs=[web_result])
        
        # 탭 5: 시스템 정보
        with gr.TabItem("⚙️ 시스템"):
            storage_btn = gr.Button("스토리지 정보 조회", variant="primary")
            storage_result = gr.Textbox(label="결과", lines=15)
            
            storage_btn.click(get_storage_info, outputs=[storage_result])


if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)