import os
import gradio as gr
import pandas as pd
import comtradeapicall
from huggingface_hub import InferenceClient
from deep_translator import GoogleTranslator

# کلید COMTRADE
subscription_key = os.getenv("COMTRADE_API_KEY", "")
# توکن Hugging Face
hf_token = os.getenv("HF_API_TOKEN")

client = InferenceClient(token=hf_token)
translator = GoogleTranslator(source='en', target='fa')

def get_importers(hs_code: str, year: str, month: str):
    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 pd.DataFrame(columns=["کد کشور", "نام کشور", "ارزش CIF"]), "برنج"
    df = df[df['cifvalue'] > 0]
    result = (
        df.groupby(["reporterCode", "reporterDesc"], as_index=False)
          .agg({"cifvalue": "sum"})
          .sort_values("cifvalue", ascending=False)
    )
    result.columns = ["کد کشور", "نام کشور", "ارزش CIF"]
    
    product_name = df['cmdDesc'].iloc[0] if 'cmdDesc' in df.columns else "برنج"
    return result, product_name

def provide_advice(table_data: pd.DataFrame, hs_code: str, year: str, month: str, product_name: str):
    if table_data is None or table_data.empty:
        return "ابتدا باید اطلاعات واردات را نمایش دهید."
    table_str = table_data.to_string(index=False)
    period = f"{year}/{int(month):02d}"
    prompt = (
        f"The following table shows countries that imported the product '{product_name}' with HS code {hs_code} during the period {period}:\n"
        f"{table_str}\n\n"
        f"لطفاً یک تحلیل کامل ارائه دهید که شامل دو بخش باشد. بخش اول فرصت‌های بازار و تقاضای بالقوه برای این محصول در این کشورها را بررسی کند، با توجه به عوامل فرهنگی، اقتصادی و جمعیتی. بخش دوم توصیه‌های استراتژیک عملی برای صادرکنندگان که این بازارها را هدف قرار داده‌اند، با تمرکز بر استراتژی‌های تجاری، مدیریت ریسک و ایجاد مشارکت‌های محلی، ارائه دهد."
    )
    print("پرامپت ساخته‌شده:")
    print(prompt)
    try:
        print("در حال فراخوانی مدل google/gemma-2b-it...")
        outputs = client.text_generation(
            prompt=prompt,
            model="google/gemma-2b-it",  # مدل سبک‌تر
            max_new_tokens=1024
        )
        print("خروجی مدل دریافت شد (به انگلیسی):")
        print(outputs)

        # ترجمه خروجی به فارسی
        translated_outputs = translator.translate(outputs)
        print("خروجی ترجمه‌شده به فارسی:")
        print(translated_outputs)
        return translated_outputs
    except Exception as e:
        error_msg = f"خطا در تولید مشاوره: {str(e)}"
        print(error_msg)
        return error_msg

current_year = pd.Timestamp.now().year
years = [str(y) for y in range(2000, current_year+1)]
months = [str(m) for m in range(1, 13)]

with gr.Blocks() as demo:
    gr.Markdown("## نمایش کشورهایی که یک کالا را وارد کرده‌اند")
    with gr.Row():
        inp_hs = gr.Textbox(label="HS Code", value="1006")
        inp_year = gr.Dropdown(choices=years, label="سال", value=str(current_year))
        inp_month = gr.Dropdown(choices=months, label="ماه", value=str(pd.Timestamp.now().month))
    btn_show = gr.Button("نمایش اطلاعات")
    out_table = gr.Dataframe(
        headers=["کد کشور", "نام کشور", "ارزش CIF"],
        datatype=["number", "text", "number"],
        interactive=True,
    )
    btn_show.click(
        fn=get_importers,
        inputs=[inp_hs, inp_year, inp_month],
        outputs=[out_table, gr.State()]
    )

    btn_advice = gr.Button("ارائه مشاوره تخصصی")
    out_advice = gr.Textbox(label="مشاوره تخصصی", lines=6)

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

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