import streamlit as st
import requests
import logging
import json
import re
from datetime import datetime, timedelta
from requests.adapters import HTTPAdapter
from urllib3.util.retry import Retry

# --- Optional: translate_query 및 관련 상수 정의 (실제 구현에 맞게 수정) ---
def translate_query(query, country):
    # 실제 상황에서는 번역 로직을 구현하세요.
    return query

COUNTRY_LOCATIONS = {
    "United States": "United States",
    "South Korea": "South Korea"
}
COUNTRY_LANGUAGES = {
    "United States": "en",
    "South Korea": "ko"
}
# -------------------------------------------------------------------------

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Page configuration
st.set_page_config(
    page_title="DeepSeek Chatbot",
    page_icon="🤖",
    layout="centered"
)

# Initialize session state for chat history
if "messages" not in st.session_state:
    st.session_state.messages = []

# Sidebar configuration
with st.sidebar:
    st.header("Model Configuration")
    st.markdown("[Get HuggingFace Token](https://huggingface.co./settings/tokens)")

    # Dropdown to select model
    model_options = [
        "deepseek-ai/DeepSeek-R1-Distill-Qwen-32B",
    ]
    selected_model = st.selectbox("Select Model", model_options, index=0)

    system_message = st.text_area(
        "System Message",
        value=(
            "You are a deep thinking AI, you may use extremely long chains of thought to deeply consider the problem and deliberate with yourself via systematic reasoning processes to help come to a correct solution prior to answering. "
            "You should enclose your thoughts and internal monologue inside tags, and then provide your solution or response to the problem."
        ),
        height=100
    )

    max_tokens = st.slider("Max Tokens", 10, 4000, 1000)
    temperature = st.slider("Temperature", 0.1, 4.0, 0.3)
    top_p = st.slider("Top-p", 0.1, 1.0, 0.6)

# Function to query the Hugging Face API
def query(payload, api_url):
    headers = {"Authorization": f"Bearer {st.secrets['HF_TOKEN']}"}
    logger.info(f"Sending request to {api_url} with payload: {payload}")
    response = requests.post(api_url, headers=headers, json=payload)
    logger.info(f"Received response: {response.status_code}, {response.text}")
    try:
        return response.json()
    except requests.exceptions.JSONDecodeError:
        logger.error(f"Failed to decode JSON response: {response.text}")
        return None

# --- 수정된 웹서치 기능 ---
def search_web(query, country="United States", page=1, num_result=10):
    url = "https://api.serphouse.com/serp/live"
    
    # 최근 24시간 기간 설정
    now = datetime.utcnow()
    yesterday = now - timedelta(days=1)
    date_range = f"{yesterday.strftime('%Y-%m-%d')},{now.strftime('%Y-%m-%d')}"
    
    # 쿼리 번역 (필요시)
    translated_query = translate_query(query, country)
    
    payload = {
        "data": {
            "q": translated_query,
            "domain": "google.com",
            "loc": COUNTRY_LOCATIONS.get(country, "United States"),
            "lang": COUNTRY_LANGUAGES.get(country, "en"),
            "device": "desktop",
            "serp_type": "web",  # 원하는 경우 "news" 등으로 변경 가능
            "page": str(page),
            "num": str(num_result),
            "date_range": date_range,
            "sort_by": "date"
        }
    }
    
    api_key = st.secrets.get("SERPHOUSE_API_TOKEN")
    if not api_key:
        logger.error("SERPHOUSE_API_TOKEN not found in st.secrets")
        return {"error": "API token not configured."}
    
    headers = {
        "accept": "application/json",
        "content-type": "application/json",
        "authorization": f"Bearer {api_key}"
    }
    
    try:
        session = requests.Session()
        retries = Retry(
            total=5,
            backoff_factor=1,
            status_forcelist=[500, 502, 503, 504, 429],
            allowed_methods=["POST"]
        )
        adapter = HTTPAdapter(max_retries=retries)
        session.mount('http://', adapter)
        session.mount('https://', adapter)
        
        response = session.post(
            url,
            json=payload,
            headers=headers,
            timeout=(30, 30)  # 연결 및 읽기 타임아웃 30초씩
        )
        response.raise_for_status()
        return {"results": response.json(), "translated_query": translated_query}
        
    except requests.exceptions.Timeout:
        return {
            "error": "검색 시간이 초과되었습니다. 잠시 후 다시 시도해주세요.",
            "translated_query": query
        }
    except requests.exceptions.RequestException as e:
        return {
            "error": f"검색 중 오류가 발생했습니다: {str(e)}",
            "translated_query": query
        }
    except Exception as e:
        return {
            "error": f"예기치 않은 오류가 발생했습니다: {str(e)}",
            "translated_query": query
        }
# --- 끝 ---

# 내부 체인 오브 쏘(생각 과정) 제거 함수
def remove_chain_of_thought(text):
    cleaned_text = re.sub(r'<think>.*?</think>', '', text, flags=re.DOTALL)
    return cleaned_text.strip()

# Chat interface
st.title("🤖 DeepSeek Chatbot")
st.caption("Powered by Hugging Face Inference API - Configure in sidebar")

# Display chat history
for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.markdown(message["content"])

# Handle input
if prompt := st.chat_input("Type your message..."):
    st.session_state.messages.append({"role": "user", "content": prompt})

    with st.chat_message("user"):
        st.markdown(prompt)

    try:
        with st.spinner("Generating response..."):
            # 업데이트된 웹서치 기능을 사용하여 검색 수행
            search_results = search_web(prompt, country="United States", page=1, num_result=10)
            
            # 검색 결과 처리
            if search_results and "results" in search_results:
                if 'organic' in search_results["results"]:
                    search_content = "\n".join(
                        [f"**{item['title']}**: {item['snippet']}" for item in search_results["results"]["organic"]]
                    )
                    search_content = f"Here are some search results related to your question:\n\n{search_content}\n\n"
                else:
                    search_content = "Sorry, no relevant search results found.\n\n"
    
                full_prompt = f"{system_message}\n\n{search_content}User: {prompt}\nAssistant:"
            else:
                full_prompt = f"{system_message}\n\nUser: {prompt}\nAssistant:"
    
            payload = {
                "inputs": full_prompt,
                "parameters": {
                    "max_new_tokens": max_tokens,
                    "temperature": temperature,
                    "top_p": top_p,
                    "return_full_text": False
                }
            }
    
            # 선택한 모델에 따른 API URL 동적 구성
            api_url = f"https://api-inference.huggingface.co/models/{selected_model}"
            logger.info(f"Selected model: {selected_model}, API URL: {api_url}")
    
            # Hugging Face API에 쿼리 전송
            output = query(payload, api_url)
    
            # API 응답 처리
            if output is not None and isinstance(output, list) and len(output) > 0:
                if 'generated_text' in output[0]:
                    assistant_response = output[0]['generated_text'].strip()
                    # 내부 체인 오브 쏘 제거
                    unique_response = remove_chain_of_thought(assistant_response)
    
                    logger.info(f"Generated response: {unique_response}")
    
                    with st.chat_message("assistant"):
                        st.markdown(unique_response)
    
                    st.session_state.messages.append({"role": "assistant", "content": unique_response})
                else:
                    logger.error(f"Unexpected API response structure: {output}")
                    st.error("Error: Unexpected response from the model. Please try again.")
            else:
                logger.error(f"Empty or invalid API response: {output}")
                st.error("Error: Unable to generate a response. Please check the model and try again.")
    
    except Exception as e:
        logger.error(f"Application Error: {str(e)}", exc_info=True)
        st.error(f"Application Error: {str(e)}")