# 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

# --- بارگذاری 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_imports(hs_code: str, year: str, month: str):
    try:
        hs_code = str(hs_code).strip()
        year = str(year).strip()
        month = str(month).strip()

        if not hs_code.isdigit() or len(hs_code) < 2:
            return "–", pd.DataFrame(), "کد HS باید یک عدد معتبر باشد (حداقل 2 رقم)!"
        if not year.isdigit() or int(year) < 2000 or int(year) > 2025:
            return "–", pd.DataFrame(), "سال باید بین 2000 و 2025 باشد!"
        if not month.isdigit() or int(month) < 1 or int(month) > 12:
            return "–", pd.DataFrame(), "ماه باید بین 1 و 12 باشد!"

        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(), "داده‌ای برای این کد HS و دوره زمانی یافت نشد."

        std_map = {
            'کد کشور': 'ptCode',
            'نام کشور': 'ptTitle',
            'ارزش CIF': 'TradeValue'
        }
        code_col = std_map['کد کشور'] if 'ptCode' in df.columns else next((c for c in df.columns if 'code' in c.lower()), None)
        title_col = std_map['نام کشور'] if 'ptTitle' in df.columns else next((c for c in df.columns if 'title' in c.lower()), None)
        value_col = std_map['ارزش CIF'] if 'TradeValue' in df.columns else 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, ""
    except Exception as e:
        return "–", pd.DataFrame(), f"خطا: {str(e)}"

# --- تابع تولید مشاوره تخصصی با 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 "ابتدا نمایش داده‌های واردات را انجام دهید."

    df_limited = table_data.head(10)
    table_str = df_limited.to_string(index=False)
    period = f"{year}/{int(month):02d}"
    prompt = (
        f"The following table shows the top {len(df_limited)} countries by CIF value importing HS code {hs_code} during {period}:\n"
        f"{table_str}\n\n"
        "Please provide a detailed and comprehensive analysis of market trends, risks, "
        "and opportunities for a new exporter entering this market."
    )
    try:
        outputs = client.text_generation(
            prompt=prompt,
            model="mistralai/Mixtral-8x7B-Instruct-v0.1",
            max_new_tokens=1024
        )
        return translator.translate(outputs)
    except Exception as e:
        return f"خطا در تولید مشاوره: {e}"

# --- CSS برای مخفی کردن هدر و فوتر و استایل بصری ---
custom_css = """
.gradio-container .header { display: none !important; }
.gradio-container .footer { display: none !important; }
.footer { display: none !important; }
.header { display: none !important; }

body {
    background-color: #1a1a1a;
    color: #e0e0e0;
}
#input-hs {
    background-color: #2c2c2c;
    color: #e0e0e0;
    border: 1px solid #444;
    border-radius: 5px;
    padding: 5px;
}
button {
    margin: 5px;
    padding: 10px 20px;
    border-radius: 5px;
}
.gr-Markdown {
    padding: 10px;
    background-color: #2c2c2c;
    border-radius: 5px;
}
.gr-Dataframe, .gr-Textbox {
    background-color: #2c2c2c;
    color: #e0e0e0;
    border: 1px solid #444;
    border-radius: 5px;
    padding: 10px;
}
"""

# --- رابط کاربری سفارشی با Gradio Blocks ---
def create_custom_interface():
    with gr.Blocks(css=custom_css) as demo:
        # عنوان برنامه
        gr.Markdown("## تحلیل واردات بر اساس کد HS و ارائه مشاوره تخصصی")

        # بخش ورودی‌ها
        with gr.Row():
            with gr.Column(scale=1):
                inp_hs = gr.Textbox(label="کد HS", placeholder="مثلاً 1006", elem_id="input-hs")
            with gr.Column(scale=1):
                inp_year = gr.Dropdown(label="سال", choices=[str(i) for i in range(2000, 2026)], value="2023")
            with gr.Column(scale=1):
                inp_month = gr.Dropdown(label="ماه", choices=[str(i).zfill(2) for i in range(1, 13)], value="01")

        # دکمه‌ها
        with gr.Row():
            btn_show = gr.Button("نمایش داده‌های واردات", variant="primary")
            btn_advice = gr.Button("ارائه مشاوره تخصصی", variant="secondary")

        # خروجی‌ها
        with gr.Row():
            out_name = gr.Markdown(label="**نام محصول**")
            out_table = gr.Dataframe(
                datatype="pandas",
                interactive=True,
                headers=["کد کشور", "نام کشور", "ارزش CIF"]
            )
            out_message = gr.Markdown(label="پیام", visible=True)

        with gr.Row():
            out_advice = gr.Textbox(label="مشاوره تخصصی", lines=8)

        # اتصال دکمه‌ها به توابع
        btn_show.click(
            fn=get_imports,
            inputs=[inp_hs, inp_year, inp_month],
            outputs=[out_name, out_table, out_message]
        )

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

    return demo

# --- راه‌اندازی رابط کاربری ---
if __name__ == "__main__":
    demo = create_custom_interface()
    demo.launch(show_api=False)  # حذف پارامتر theme