import os
import zipfile
import json
from flask import Flask, request, abort
from linebot import LineBotApi, WebhookHandler
from linebot.exceptions import InvalidSignatureError
from linebot.models import MessageEvent, TextMessage, TextSendMessage

# RAG & AI 相關
import aisuite as ai
from langchain_community.vectorstores import FAISS
from langchain_community.embeddings import HuggingFaceEmbeddings
from huggingface_hub import login

# ==========================================
# 1. 初始化與環境設定
# ==========================================
app = Flask(__name__)

# 讀取 Secrets
GROQ_API_KEY = os.environ.get('GROQ')
HF_TOKEN = os.environ.get('HuggingFace')
LINE_CHANNEL_ACCESS_TOKEN = os.environ.get('LINE_CHANNEL_ACCESS_TOKEN')
LINE_CHANNEL_SECRET = os.environ.get('LINE_CHANNEL_SECRET')

# 設定環境變數供 aisuite 使用
if GROQ_API_KEY:
    os.environ['GROQ_API_KEY'] = GROQ_API_KEY

# 登入 Hugging Face
if HF_TOKEN:
    try:
        login(token=HF_TOKEN)
        print("✅ HF Login Success")
    except Exception as e:
        print(f"❌ HF Login Failed: {e}")

# 初始化 LINE API
line_bot_api = LineBotApi(LINE_CHANNEL_ACCESS_TOKEN)
handler = WebhookHandler(LINE_CHANNEL_SECRET)

# ==========================================
# 2. RAG 系統初始化
# ==========================================
# 解壓縮資料庫
if not os.path.exists("faiss_dbV2"):
    if os.path.exists("faiss_dbV2.zip"):
        print("Unzipping database...")
        with zipfile.ZipFile("faiss_dbV2.zip", 'r') as zip_ref:
            zip_ref.extractall(".")
    else:
        print("⚠️ Warning: faiss_dbV2.zip not found.")

# Embedding 模型
class EmbeddingGemmaEmbeddings(HuggingFaceEmbeddings):
    def __init__(self, **kwargs):
        super().__init__(model_name="google/embeddinggemma-300m", encode_kwargs={"normalize_embeddings": True}, **kwargs)
    def embed_documents(self, texts):
        return super().embed_documents([f"title: none | text: {t}" for t in texts])
    def embed_query(self, text):
        return super().embed_query(f"task: search result | query: {text}")

print("Loading FAISS...")
embedding_model = EmbeddingGemmaEmbeddings()
try:
    vectorstore = FAISS.load_local("faiss_db", embeddings=embedding_model, allow_dangerous_deserialization=True)
    retriever = vectorstore.as_retriever(search_kwargs={"k": 5}) # k 值可以稍微調小，因為純文字不需要太多來源
    print("✅ RAG System Ready")
except Exception as e:
    print(f"❌ RAG Init Failed: {e}")
    retriever = None

# AI Client
client = ai.Client()

# ==========================================
# 3. Prompt 設定 (核心修改處：全文字模式)
# ==========================================
PROMPTS = {
    "planner": r"""
    你是一位實變函數論專家。
    使用者會問關於數學的問題。
    請分析問題的核心邏輯，提出一種邏輯清晰、正確的思路。
    """,
    
    "writer": r"""
    你是一位非常嚴謹的的數學教授，擅長用白話文解釋數學。
    請參考 Planner 提供的思路，引用相關的定義與定理 (Definition/Theorem)，撰寫完整的解答： 
    【絕對規則】：
    1. **禁止使用任何 LaTeX 代碼**（絕對不要出現 $符號、\int、\infty 這種東西）。
    2. **禁止使用數學符號**。請用中文口語代替。
       - ❌ 錯誤：$f(x) > 0$
       - ✅ 正確：函數值大於零
       - ❌ 錯誤：$x \in A$
       - ✅ 正確：x 屬於集合 A
    3. 請用繁體中文回答。
    4. 結構要清晰：先講結論（有/沒有），再講原因，最後舉個清楚、正確的例子。
    """,
    
    "reviewer": r"""
    你是一位負責嚴格數學審稿人兼文字編輯。
    請檢查這份數學解答。
    檢查重點：
    1. 定義引用是否正確？
    2. 證明邏輯是否有漏洞 (Logical Gaps)？
    3. 是否有濫用直觀而犧牲嚴謹性的情況？
    4. 檢查文章中**是否混入了數學符號或 LaTeX 代碼**。如果有，請把它們全部改成中文敘述。
    並請給出具體的修改建議。
    如果解答正確且文章通順且完全是純文字，請回答 PASS。
    """,
    
    "refiner": r"""
    你是最終定稿人。
    請根據建議，修正並潤飾最終的數學解答。
    請保留原本正確的部分，針對被指出的錯誤進行修正。
    確保整篇回答看起來就像是 LINE 上面朋友的對話，親切且沒有閱讀門檻。
    """
}

# ==========================================
# 4. 核心邏輯函式
# ==========================================
def call_ai(system, user, model_id="groq:llama-3.3-70b-versatile"):
    try:
        resp = client.chat.completions.create(
            model=model_id,
            messages=[{"role": "system", "content": system}, {"role": "user", "content": user}]
        )
        return resp.choices[0].message.content
    except Exception as e:
        return f"Error: {e}"

def math_solver(user_question):
    if not retriever: return "系統資料庫未載入，無法回答問題。"
    
    # 1. 檢索
    docs = retriever.invoke(user_question)
    context = "\n".join([d.page_content for d in docs])
    
    # 2. 規劃 (Planner)
    plan = call_ai(PROMPTS["planner"], f"問題: {user_question}\n背景知識: {context}")
    
    # 3. 撰寫 (Writer) - 
    draft = call_ai(PROMPTS["writer"], f"問題: {user_question}\n寫作策略: {plan}\n背景知識: {context}")
    
    # 4. 審查 (Reviewer) - 確保沒有漏網的 LaTeX
    review = call_ai(PROMPTS["reviewer"], f"草稿: {draft}")
    
    if "PASS" not in review:
        final = call_ai(PROMPTS["refiner"], f"草稿: {draft}\n修改建議: {review}")
    else:
        final = draft
        
    return final

# ==========================================
# 5. Flask & LINE Webhook 路由
# ==========================================
@app.route("/")
def home():
    return "Text-Only Math AI Agent is Running!"

@app.route("/callback", methods=['POST'])
def callback():
    signature = request.headers['X-Line-Signature']
    body = request.get_data(as_text=True)
    
    try:
        handler.handle(body, signature)
    except InvalidSignatureError:
        abort(400)
    return 'OK'

@handler.add(MessageEvent, message=TextMessage)
def handle_message(event):
    user_msg = event.message.text
    
    try:
        # 1. 執行 AI 邏輯
        answer = math_solver(user_msg)
        
        # 2. 直接回傳純文字
        line_bot_api.reply_message(
            event.reply_token,
            TextSendMessage(text=answer)
        )
        
    except Exception as e:
        error_msg = f"系統忙碌中，請稍後再試。\n錯誤: {str(e)}"
        print(error_msg)
        line_bot_api.reply_message(event.reply_token, TextSendMessage(text=error_msg))

# ==========================================
# 6. 啟動
# ==========================================
if __name__ == "__main__":
    app.run(host="0.0.0.0", port=7860)