# app.py
import os
import pandas as pd
import gradio as gr
import comtradeapicall
from huggingface_hub import InferenceClient
from deep_translator import GoogleTranslator
import spaces  # برای مدیریت GPU کرایه‌ای

# --- بارگذاری HS DATA از CSV گیت‌هاب ---
HS_CSV_URL = (
    "https://raw.githubusercontent.com/"
    "datasets/harmonized-system/master/data/harmonized-system.csv"
)
hs_df = pd.read_csv(HS_CSV_URL, dtype=str)

def get_product_name(hs_code: str) -> str:
    code4 = str(hs_code).zfill(4)
    row = hs_df[hs_df["hscode"] == code4]
    return row.iloc[0]["description"] if not row.empty else "–"

# --- تابع دریافت واردات و پردازش ستون‌ها ---
def get_importers(hs_code: str, year: str, month: str):
    product_name = get_product_name(hs_code)
    period = f"{year}{int(month):02d}"
    df = comtradeapicall.previewFinalData(
        typeCode='C', freqCode='M', clCode='HS', period=period,
        reporterCode=None, cmdCode=hs_code, flowCode='M',
        partnerCode=None, partner2Code=None,
        customsCode=None, motCode=None,
        maxRecords=500, includeDesc=True
    )
    if df is None or df.empty:
        return product_name, pd.DataFrame()

    # شناسایی ستون‌های مورد نیاز
    code_col = next((c for c in df.columns if 'code' in c.lower()), None)
    title_col = next((c for c in df.columns if 'title' in c.lower()), None)
    value_col = next((c for c in df.columns if 'value' in c.lower()), None)
    if not (code_col and title_col and value_col):
        return product_name, df

    df_sorted = df.sort_values(value_col, ascending=False).head(10)
    out = df_sorted[[code_col, title_col, value_col]]
    out.columns = ['کد کشور', 'نام کشور', 'ارزش CIF']
    return product_name, out

# --- تابع تولید مشاوره تخصصی با GPU کرایه‌ای ---
hf_token = os.getenv("HF_API_TOKEN")
client = InferenceClient(token=hf_token)
translator = GoogleTranslator(source='en', target='fa')

@spaces.GPU
def provide_advice(table_data: pd.DataFrame, hs_code: str, year: str, month: str):
    if table_data is None or table_data.empty:
        return "<div dir=\"rtl\" style=\"text-align: right; font-family: IRANSans;\">ابتدا نمایش داده‌های واردات را انجام دهید.</div>"

    df_limited = table_data.head(10)
    table_str = df_limited.to_string(index=False)
    period = f"{year}/{int(month):02d}"
    prompt = f"""
لطفاً یک گزارش تحلیلی به زبان فارسی و با چینش راست‌به‌چپ تولید کن، شامل دو بخش:
1. جدول خلاصه (Markdown) 10 کشور برتر با ستون‌های رتبه، نام کشور، ارزش CIF و سه تأمین‌کننده برتر.
2. نکات کلیدی و پیشنهادات به‌صورت بولت‌پوینت با تراز راست.

- متن خروجی را درون یک <div dir=\"rtl\" style=\"text-align: right; font-family: IRANSans;\"> قرار بده.

داده‌ها:
{table_str}
""" 10 کشور برتر با ستون‌های رتبه، نام کشور، ارزش CIF و سه تأمین‌کننده برتر."
        "\n2. نکات کلیدی و پیشنهادات به‌صورت بولت‌پوینت با تراز راست."
        "\n
- متن خروجی را درون یک <div dir=\\\"rtl\\\" style=\\\"text-align: right; font-family: IRANSans;\\\"> قرار بده."
        f"\n
داده‌ها:\n{table_str}"
    )
    try:
        outputs = client.text_generation(
            prompt=prompt,
            model="mistralai/Mixtral-8x7B-Instruct-v0.1",
            max_new_tokens=2048
        )
        persian = translator.translate(outputs)
        # قرار دادن در کانتینر HTML راست‌چین
        return f"<div dir=\"rtl\" style=\"text-align: right; font-family: IRANSans;\">{persian}</div>"
    except Exception as e:
        return f"<div dir=\"rtl\" style=\"text-align: right; font-family: IRANSans;\">خطا در تولید مشاوره: {e}</div>"

# --- رابط کاربری Gradio ---
with gr.Blocks() as demo:
    gr.Markdown("<div dir=\"rtl\" style=\"text-align: right; font-family: IRANSans;\"><h2>تحلیل واردات بر اساس کد HS و ارائه مشاوره تخصصی</h2></div>", elem_id="title")

    with gr.Row():
        inp_hs = gr.Textbox(label="کد HS", placeholder="مثلاً 1006")
        inp_year = gr.Textbox(label="سال", placeholder="مثلاً 2023")
        inp_month = gr.Textbox(label="ماه", placeholder="مثلاً 1 تا 12")

    btn_show = gr.Button("نمایش داده‌های واردات")
    out_name = gr.Markdown(label="**<div dir=\"rtl\" style=\"text-align: right; font-family: IRANSans;\">نام محصول</div>")
    out_table = gr.Dataframe(
        datatype="pandas",
        interactive=True
    )

    btn_show.click(
        fn=get_importers,
        inputs=[inp_hs, inp_year, inp_month],
        outputs=[out_name, out_table]
    )

    btn_advice = gr.Button("ارائه مشاوره تخصصی")
    out_advice = gr.HTML()

    btn_advice.click(
        fn=provide_advice,
        inputs=[out_table, inp_hs, inp_year, inp_month],
        outputs=out_advice
    )

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