import os
import gradio as gr
import pandas as pd
import comtradeapicall
from huggingface_hub import InferenceClient
from googletrans import Translator  # اضافه کردن کتابخانه ترجمه

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

client = InferenceClient(token=hf_token)
translator = Translator()  # ایجاد یک نمونه از مترجم

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"]
    return result

def provide_advice(table_data: pd.DataFrame, hs_code: str, year: str, month: 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 a product with HS code {hs_code} during the period {period}:\n"
        f"{table_str}\n\n"
        f"Please provide a detailed analysis in two complete paragraphs. The first paragraph should discuss the market opportunities and potential demand for this product in these countries. The second paragraph should offer strategic recommendations for exporters targeting these markets, focusing on trade strategies and risk management."
    )
    print("پرامپت ساخته‌شده:")
    print(prompt)
    try:
        print("در حال فراخوانی مدل mistralai/Mixtral-8x7B-Instruct-v0.1...")
        outputs = client.text_generation(
            prompt=prompt,
            model="mistralai/Mixtral-8x7B-Instruct-v0.1",
            max_new_tokens=512
        )
        print("خروجی مدل دریافت شد (به انگلیسی):")
        print(outputs)

        # ترجمه خروجی به فارسی
        translated_outputs = translator.translate(outputs, src='en', dest='fa').text
        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")
        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(get_importers, [inp_hs, inp_year, inp_month], out_table)

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

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

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