from groq import Groq
import gradio as gr
import os

client = Groq(api_key=os.getenv("GROQ_API_KEY"))

def get_chatbot_response(user_message, country, language, conversation_history):

    system_message = (
        f"You are a lawyer specializing in providing concise and accurate legal information based on the laws in {country}. "
        f"Respond in {language}. Provide clear, factual information without offering personal legal advice or opinions. "
        "Include relevant legal references, statutes, or case law when possible."
    )

    if conversation_history:
        if conversation_history[0]["role"] == "system":
            conversation_history[0]["content"] = system_message
        else:
            conversation_history.insert(0, {"role": "system", "content": system_message})
    else:
        conversation_history.append({"role": "system", "content": system_message})
    
    conversation_history.append({"role": "user", "content": user_message})

    completion = client.chat.completions.create(
        model="deepseek-r1-distill-llama-70b",
        messages=conversation_history,
        temperature=0.3,
        top_p=0.95,
        stream=True,
        reasoning_format="hidden"
    )

    response = ""
    for chunk in completion:
        response += chunk.choices[0].delta.content or ""

    conversation_history.append({"role": "assistant", "content": response})

    chat_display = [
        (msg["content"], conversation_history[i + 1]["content"]) 
        for i, msg in enumerate(conversation_history[:-1]) if msg["role"] == "user"
    ]
    
    return conversation_history, chat_display

theme = gr.themes.Ocean(
    text_size="lg",
    font=[gr.themes.GoogleFont('DM Sans'), 'ui-sans-serif', 'system-ui', 'sans-serif'],
).set(
    body_text_size='*text_lg',
    background_fill_secondary='*secondary_100',
    chatbot_text_size='*text_lg',
    input_radius='*radius_md',
    input_text_size='*text_lg',
)

custom_css = """
.title-text {
    background: #00A0B0;
    -webkit-background-clip: text;
    background-clip: text;
    color: transparent;
    -webkit-text-fill-color: transparent;
    display: inline-block;
    width: fit-content;
    font-weight: bold;
    text-align: center;
    font-size: 45px;
}

.law-button {
    border: 1px solid #00A0B0;
    background-color: transparent;
    font-size: 15px;
    padding: 5px 15px;
    border-radius: 16px;
    margin: 0 5px;
}

.law-button:hover {
    background: linear-gradient(90deg, #00A0B0, #00FFEF);
    color: white;
}
"""

def clear_history():
    return []


with gr.Blocks(theme = theme, css = custom_css) as demo:
    gr.HTML("<h2 class='title-text'>⚖️ AI Legal Chatbot</h2>")
    gr.Markdown("### Hey there! Pick your country, choose a language, and tell us about your legal situation. We're here to help!")

    with gr.Row():
        country_input = gr.Dropdown(
            ["Canada", "United States", "United Kingdom", "Spain", "France", "Germany", "India", "China", "Lebanon", "Other"],
            label="🌍 Select Country",
            interactive=True
        )
        language_input = gr.Dropdown(
            ["English", "Spanish", "French", "German", "Hindi", "Mandarin", "Arabic"],
            label="🗣️ Select Language",
            interactive=True
        )

    custom_country_input = gr.Textbox(label="Enter Country (if not listed)", visible=False)

    conversation_state = gr.State([])

    chatbot = gr.Chatbot(label="💬 Chat History")
    chatbot.clear(fn=clear_history, outputs=conversation_state)
    
    with gr.Row():
        family_btn = gr.Button("👨‍👩‍👧‍👦 Family", elem_classes="law-button")
        corporate_btn = gr.Button("🏢 Corporate", elem_classes="law-button")
        health_btn = gr.Button("🏥 Health", elem_classes="law-button")
        military_btn = gr.Button("🎖️ Military", elem_classes="law-button")
        immigration_btn = gr.Button("🌍 Immigration", elem_classes="law-button")
        criminal_btn = gr.Button("🚔 Criminal", elem_classes="law-button")
        property_btn = gr.Button("🏠 Property", elem_classes="law-button")
        environmental_btn = gr.Button("🌱 Environmental", elem_classes="law-button")


    scenario_input = gr.Textbox(label="💡 Type your message...", placeholder="Describe your legal situation...", interactive=True)

    def update_law_selection(current, new_selection):
        if "Law:" in current:
            parts = current.split("Law:", 1)
            additional_text = parts[1] if len(parts) > 1 else ""
        else:
            additional_text = current

        return f"{new_selection} Law: {additional_text}"

    for btn, law in zip(
        [family_btn, corporate_btn, health_btn, military_btn, immigration_btn, criminal_btn, property_btn, environmental_btn],
        ["Family", "Corporate", "Health", "Military", "Immigration", "Criminal", "Property", "Environmental"]
    ):
        btn.click(lambda current, law=law: update_law_selection(current, law), inputs=scenario_input, outputs=scenario_input)

    submit_btn = gr.Button("Send",variant="primary")

    def submit(country, custom_country, language, scenario, conversation_state):
        selected_country = custom_country if country == "Other" else country
        updated_history, chat_display = get_chatbot_response(
            scenario, selected_country, language, conversation_state
        )
        return updated_history, chat_display, ""


    country_input.change(lambda c: gr.update(visible=c == "Other"), inputs=country_input, outputs=custom_country_input)

    submit_btn.click(
        submit, 
        inputs=[country_input, custom_country_input, language_input, scenario_input, conversation_state], 
        outputs=[conversation_state, chatbot, scenario_input]
    )

demo.launch()