import os
import re
import tempfile
from pathlib import Path
from typing import Tuple

import gradio as gr
import pandas as pd
from openai import OpenAI

# reportlab for PDF
from reportlab.lib import colors
from reportlab.lib.pagesizes import A4
from reportlab.platypus import SimpleDocTemplate, Table, TableStyle, Paragraph, Spacer
from reportlab.lib.styles import getSampleStyleSheet, ParagraphStyle
from reportlab.pdfbase import pdfmetrics
from reportlab.pdfbase.cidfonts import UnicodeCIDFont

# --------------------
# Setup OpenAI client
# --------------------
API_KEY = os.getenv("OPENAI_API_KEY") or os.getenv("OPENAIAPIKEY")
client = OpenAI(api_key=API_KEY) if API_KEY else None
FALLBACK_TEXT = """
### 1. 可能相關疾病
- 胃食道逆流（GERD）：胃酸回流到食道，可能引起燒灼感和胃痛。
- 胃潰瘍：胃內部的黏膜受損，形成潰瘍，可能導致疼痛和不適。
### 2. 飲食營養建議
- 纖維素：有助於消化，減少便秘和腸道不適。
- 維生素C：幫助修復胃黏膜，增強免疫系統。
- Omega-3脂肪酸：具有抗炎特性，可能有助於減輕胃部炎症。
### 3. 一日三餐具體菜單
#### 中式版本
- 早餐: 燕麥粥（牛奶）、水煮蛋
- 午餐: 清蒸魚、炒青菜、白飯
- 晚餐: 魚湯、燙青菜、糙米
#### 西式版本
- 早餐: 全麥吐司、水煮蛋、藍莓
- 午餐: 烤雞胸肉、藜麥沙拉、優格
- 晚餐: 煎三文魚、蒸花椰菜、藜麥
""".strip()

# --------------------
# GPT 呼叫（加速版）
# --------------------
def call_gpt_system(prompt: str) -> str:
    if client is None:
        return FALLBACK_TEXT
    try:
        resp = client.chat.completions.create(
            model="gpt-4o-mini",
            max_tokens=500,         # 限制字數以加快速度
            temperature=0.6,        # 適度隨機
            messages=[{"role": "user", "content": prompt}],
        )
        text = resp.choices[0].message.content
        return text if text and text.strip() else FALLBACK_TEXT
    except Exception as e:
        print("GPT 呼叫失敗:", e)
        return FALLBACK_TEXT

# --------------------
# 區塊解析
# --------------------
def extract_block(text: str, starts: list, ends: list = None) -> str:
    if not text:
        return ""
    start_pos = None
    for s in starts:
        m = re.search(s, text, flags=re.I)
        if m:
            start_pos = m.end()
            break
    if start_pos is None:
        return ""
    if not ends:
        return text[start_pos:].strip()
    end_pos = len(text)
    for e in ends:
        m = re.search(e, text[start_pos:], flags=re.I)
        if m:
            end_pos = start_pos + m.start()
            break
    return text[start_pos:end_pos].strip()

def parse_bullets_to_pairs(block: str):
    items = []
    if not block:
        return items
    lines = [ln.strip() for ln in block.splitlines() if ln.strip()]
    candidate = [ln for ln in lines if ln.lstrip().startswith("-")] or lines
    for ln in candidate:
        ln2 = re.sub(r"^-+\s*", "", ln).strip()
        if "：" in ln2:
            name, desc = ln2.split("：", 1)
            items.append((name.strip(), desc.strip()))
        elif ":" in ln2:
            name, desc = ln2.split(":", 1)
            items.append((name.strip(), desc.strip()))
        else:
            items.append((ln2.strip(), ""))
    return items

def parse_menu_block(block: str):
    rows = []
    if not block:
        return rows
    lines = [ln.strip() for ln in block.splitlines() if ln.strip()]
    for ln in lines:
        ln2 = re.sub(r"^-+\s*", "", ln).strip()
        if "：" in ln2:
            meal, food = ln2.split("：", 1)
            rows.append({"meal": meal.strip(), "food": food.strip()})
        elif ":" in ln2:
            meal, food = ln2.split(":", 1)
            rows.append({"meal": meal.strip(), "food": food.strip()})
        else:
            rows.append({"meal": "", "food": ln2})
    return rows

def robust_parse(text: str) -> Tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame]:
    disease_block = extract_block(text, [r"可能相關疾病", r"###\s*1", r"1\."], [r"飲食", r"###\s*2", r"2\."])
    nutrition_block = extract_block(text, [r"飲食營養建議", r"###\s*2", r"2\."], [r"一日三餐", r"###\s*3", r"3\."])
    menu_block = extract_block(text, [r"一日三餐", r"###\s*3", r"3\."], [r"參考來源", r"$"])

    diseases = parse_bullets_to_pairs(disease_block)
    nutritions = parse_bullets_to_pairs(nutrition_block)

    ch_block = ""
    we_block = ""
    m_ch = re.search(r"(?:中式版本|中式)\s*(.*?)(?=(?:西式版本|西式|$))", menu_block, flags=re.S | re.I)
    if m_ch:
        ch_block = m_ch.group(1)
        m_we = re.search(r"(?:西式版本|西式)\s*(.*)", menu_block[m_ch.end():], flags=re.S | re.I)
        if m_we:
            we_block = m_we.group(1)
    else:
        ch_block = menu_block

    ch_rows = parse_menu_block(ch_block)
    we_rows = parse_menu_block(we_block)

    if not diseases:
        diseases = parse_bullets_to_pairs(extract_block(FALLBACK_TEXT, [r"可能相關疾病"], None))
    if not nutritions:
        nutritions = parse_bullets_to_pairs(extract_block(FALLBACK_TEXT, [r"飲食營養建議"], None))

    df_disease = pd.DataFrame([{"疾病": n, "說明": d} for n, d in diseases]) if diseases else pd.DataFrame(columns=["疾病", "說明"])
    df_nutrition = pd.DataFrame([{"營養素": n, "建議": d} for n, d in nutritions]) if nutritions else pd.DataFrame(columns=["營養素", "建議"])

    max_len = max(len(ch_rows), len(we_rows))
    meals = []
    for i in range(max_len):
        ch = ch_rows[i]["food"] if i < len(ch_rows) else ""
        we = we_rows[i]["food"] if i < len(we_rows) else ""
        meals.append({"中式餐點": ch, "西式餐點": we})
    df_meals = pd.DataFrame(meals, columns=["中式餐點", "西式餐點"])
    return df_disease, df_nutrition, df_meals

# --------------------
# PDF 生成（修正版：文字換行）
# --------------------
def has_chinese(s: str) -> bool:
    return bool(re.search(r"[\u4e00-\u9fff]", str(s)))

def clean_dataframe(df: pd.DataFrame) -> pd.DataFrame:
    if df.empty:
        return df
    mask = df.apply(lambda row: any(has_chinese(x) for x in row), axis=1)
    return df[mask].reset_index(drop=True)

def generate_pdf_file(df_disease, df_nutrition, df_meals, basic_info=None):
    df_disease = clean_dataframe(df_disease)
    df_nutrition = clean_dataframe(df_nutrition)
    df_meals = clean_dataframe(df_meals)

    pdfmetrics.registerFont(UnicodeCIDFont("STSong-Light"))
    font_name = "STSong-Light"

    file_prefix = basic_info.get("name", "個人") if basic_info else "個人"
    safe_name = re.sub(r"[^\w\u4e00-\u9fff]", "_", file_prefix)
    filename = f"{safe_name}的專屬食療報告書.pdf"

    tmp_path = Path(tempfile.gettempdir()) / filename
    doc = SimpleDocTemplate(str(tmp_path), pagesize=A4, leftMargin=36, rightMargin=36, topMargin=36, bottomMargin=36)
    elements = []
    styles = getSampleStyleSheet()
    title_style = styles["Title"].clone("title")
    heading_style = styles["Heading2"].clone("h2")
    normal_style = ParagraphStyle("normal", fontName=font_name, fontSize=10, leading=14)

    title_style.fontName = heading_style.fontName = font_name

    # 標題
    elements.append(Paragraph(f"健康分析與飲食建議報告 — {file_prefix}", title_style))
    elements.append(Spacer(1, 12))

    # 基本資料摘要表格
    if basic_info:
        elements.append(Paragraph("基本資料摘要", heading_style))
        table_data = [
            ["姓名", basic_info.get("name", "")],
            ["性別", basic_info.get("gender", "")],
            ["年齡", str(basic_info.get("age", ""))],
            ["身高", f"{basic_info.get('height', '')} cm"],
            ["體重", f"{basic_info.get('weight', '')} kg"],
        ]
        t = Table(table_data, colWidths=[80, 400])
        style = TableStyle([
            ("GRID", (0, 0), (-1, -1), 0.8, colors.black),
            ("BACKGROUND", (0, 0), (-1, 0), colors.lightgrey),
            ("FONTNAME", (0, 0), (-1, -1), font_name),
        ])
        t.setStyle(style)
        elements.append(t)
        elements.append(Spacer(1, 16))

    # 疾病表
    if not df_disease.empty:
        elements.append(Paragraph("一、可能相關疾病", heading_style))
        table_data = [df_disease.columns.tolist()] + [[Paragraph(str(c), normal_style) for c in row] for row in df_disease.values.tolist()]
        t = Table(table_data, colWidths=[150, 330])
        t.setStyle(TableStyle([("GRID", (0, 0), (-1, -1), 0.8, colors.black), ("BACKGROUND", (0, 0), (-1, 0), colors.lightgrey), ("FONTNAME", (0, 0), (-1, -1), font_name)]))
        elements.append(t)
        elements.append(Spacer(1, 12))

    # 營養建議表
    if not df_nutrition.empty:
        elements.append(Paragraph("二、飲食營養建議", heading_style))
        table_data = [df_nutrition.columns.tolist()] + [[Paragraph(str(c), normal_style) for c in row] for row in df_nutrition.values.tolist()]
        t = Table(table_data, colWidths=[150, 330])
        t.setStyle(TableStyle([("GRID", (0, 0), (-1, -1), 0.8, colors.black), ("BACKGROUND", (0, 0), (-1, 0), colors.lightgrey), ("FONTNAME", (0, 0), (-1, -1), font_name)]))
        elements.append(t)
        elements.append(Spacer(1, 12))

    # 餐點表（換行修正）
    if not df_meals.empty:
        elements.append(Paragraph("三、一日三餐（中式 / 西式）", heading_style))
        table_data = [df_meals.columns.tolist()] + [[Paragraph(str(c), normal_style) for c in row] for row in df_meals.values.tolist()]
        t = Table(table_data, colWidths=[240, 240])
        t.setStyle(TableStyle([("GRID", (0, 0), (-1, -1), 0.8, colors.black), ("BACKGROUND", (0, 0), (-1, 0), colors.lightgrey), ("FONTNAME", (0, 0), (-1, -1), font_name)]))
        elements.append(t)

    doc.build(elements)
    return str(tmp_path)

# --------------------
# 主處理函式
# --------------------
def on_submit(name, height, weight, age, gender, symptom):
    prompt = f"""你是一位專業營養師。根據以下使用者基本資料與症狀，請給出：
1) 可能相關疾病（每行 - 疾病：說明）
2) 飲食營養建議（每行 - 營養素：建議）
3) 一日三餐具體菜單（分中式版本與西式版本，每項用 - 早餐: ... 格式）
使用者資料：
- 姓名: {name}
- 性別: {gender}
- 年齡: {age}
- 身高: {height} cm
- 體重: {weight} kg
症狀: {symptom}
請用中文輸出，格式如上，方便解析。"""
    raw_text = call_gpt_system(prompt)
    df_disease, df_nutrition, df_meals = robust_parse(raw_text)
    basic_info = {"name": name, "gender": gender, "age": age, "height": height, "weight": weight}
    pdf_path = generate_pdf_file(df_disease, df_nutrition, df_meals, basic_info)
    return df_disease, df_nutrition, df_meals, gr.update(value=pdf_path, visible=True)

# --------------------
# Gradio UI + CSS 樣式
# --------------------
css_code = """
#submit_btn {background-color: #007bff !important; color: white !important;}
#download_btn {background-color: orange !important; color: white !important;}
"""

with gr.Blocks(css=css_code) as demo:
    gr.Markdown("## 🩺AI食療專家")
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### 請輸入基本資料")
            name = gr.Textbox(label="姓名", placeholder="請輸入姓名")
            gender = gr.Dropdown(["男", "女"], label="性別", value="男")
            age = gr.Number(label="年齡", value=30, precision=0)
            height = gr.Number(label="身高 (cm)", value=170, precision=0)
            weight = gr.Number(label="體重 (kg)", value=65, precision=0)
            symptom = gr.Textbox(label="症狀 / 主訴", placeholder="例如：胃痛、頭痛")
            submit_btn = gr.Button("分析並生成報告", elem_id="submit_btn")
        with gr.Column(scale=2):
            gr.Markdown("### 判讀結果表格")
            df_disease_out = gr.Dataframe(label="可能相關疾病", interactive=False)
            df_nutrition_out = gr.Dataframe(label="飲食營養建議", interactive=False)
            df_meals_out = gr.Dataframe(label="中式 / 西式 餐點對照", interactive=False)
            download_btn = gr.DownloadButton(label="一鍵下載 PDF", visible=False, elem_id="download_btn")

    submit_btn.click(
        fn=on_submit,
        inputs=[name, height, weight, age, gender, symptom],
        outputs=[df_disease_out, df_nutrition_out, df_meals_out, download_btn],
    )

if __name__ == "__main__":
    demo.launch()