import os
import gradio as gr
import pandas as pd
import comtradeapicall
from huggingface_hub import InferenceApi

subscription_key = os.getenv("COMTRADE_API_KEY", "")
hf_token         = os.getenv("HF_API_TOKEN", None)
proxy_url        = None

inference = InferenceApi(repo_id="google/gemma-2b", token=hf_token)

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,
        proxy_url=proxy_url
    )
    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: pd.DataFrame, hs_code: str, year: str, month: str):
    if table is None or table.empty:
        return "ابتدا باید اطلاعات واردات را نمایش دهید."
    table_md = table.to_markdown(index=False)
    period = f"{year}/{int(month):02d}"
    prompt = (
        f"جدول زیر کشورهایی را نشان می‌دهد که کالا با کد HS {hs_code} را در دوره {period} وارد کرده‌اند:\n"
        f"{table_md}\n\n"
        "لطفاً بر اساس این اطلاعات دو پاراگراف مشاوره تخصصی بنویسید."
    )
    response = inference(prompt)
    return response

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=False,
    )
    btn_show.click(get_importers, inputs=[inp_hs, inp_year, inp_month], outputs=out_table)

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

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

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