import gradio as gr
import os
from huggingface_hub import InferenceClient

# Setup HF Token
token_path = os.path.expanduser("~/.cache/huggingface/token")
HF_TOKEN = os.environ.get("HF_TOKEN")
if not HF_TOKEN and os.path.exists(token_path):
    with open(token_path) as f:
        HF_TOKEN = f.read().strip()

# Model Config - Using the STABLE base model for reliable Cloud Inference
MODEL_ID = "Qwen/Qwen2.5-1.5B-Instruct" 
client = InferenceClient(model=MODEL_ID, token=HF_TOKEN)

# THE REAL SYSTEM PROMPT
system_prompt = """You are LegalBuddy, a professional legal document drafting assistant for Indian law. 
Your objective is to help users generate highly accurate, structured legal documents.

STRICT INSTRUCTIONS:
1. INITIAL LANGUAGE: Always start in English. 
2. DYNAMIC LANGUAGE: If the user speaks in Hindi/Hinglish, you MUST respond in the same. Otherwise, stick to English.
3. INTERVIEW MODE: Ask structured questions ONE AT A TIME to collect missing info (Landlord, Tenant, Rent, etc.).
4. DRAFTING: When ready, generate the full professional legal document structure with # Headers and clear clauses.
"""

custom_css = """
body, .gradio-container { font-family: 'Inter', -apple-system, sans-serif !important; background-color: #f8fafc !important; }
#header { padding: 30px; background: linear-gradient(135deg, #1e293b 0%, #0f172a 100%); border-radius: 12px; margin-bottom: 25px; box-shadow: 0 4px 6px -1px rgba(0,0,0,0.1); text-align: center; }
#header h1 { margin: 0; font-size: 32px; font-weight: 800; color: #ffffff !important; letter-spacing: -0.5px; }
#header p { margin: 8px 0 0 0; font-size: 16px; color: #cbd5e1 !important; font-weight: 400; }
.chatbot-container { border-radius: 12px !important; box-shadow: 0 10px 25px -5px rgba(0, 0, 0, 0.1) !important; background: white !important; }
.message-wrap { font-size: 16px !important; line-height: 1.6 !important; }
"""

def setup_chat(user_text, history):
    history.append((user_text, ""))
    return gr.update(value="", interactive=False), history, gr.update(visible=False), gr.update(visible=True)

def chat_logic(history, temp, top_p_val, max_tokens):
    messages = [{"role": "system", "content": system_prompt}]
    for u_msg, a_reply in history[:-1]:
        if u_msg: messages.append({"role": "user", "content": u_msg})
        if a_reply: messages.append({"role": "assistant", "content": a_reply})
    
    messages.append({"role": "user", "content": history[-1][0]})
    
    partial_response = ""
    try:
        response_stream = client.chat_completion(
            messages,
            max_tokens=int(max_tokens),
            stream=True,
            temperature=float(temp),
            top_p=float(top_p_val),
        )
        for chunk in response_stream:
            if chunk.choices and chunk.choices[0].delta.content:
                partial_response += chunk.choices[0].delta.content
                yield partial_response
    except Exception as e:
        yield f"⚠️ Connection Issue: {str(e)}"

def process_interaction(chat_history, temp, top_p_val, max_tokens):
    user_input = chat_history[-1][0]
    for partial_response in chat_logic(chat_history, temp, top_p_val, max_tokens):
        chat_history[-1] = (user_input, partial_response)
        yield chat_history

def finalize_chat():
    return gr.update(interactive=True), gr.update(visible=True), gr.update(visible=False)

with gr.Blocks(theme=gr.themes.Default(primary_hue="slate", neutral_hue="slate"), css=custom_css, title="LegalBuddy Pro") as demo:
    with gr.Column(elem_id="header"):
        gr.Markdown("<h1>LegalBuddy Pro</h1>\n<p>Professional Legal Drafting Assistant</p>")
    
    with gr.Row():
        with gr.Column(scale=12): # Full Width
            chatbot = gr.Chatbot(
                height=650, 
                show_label=False, 
                show_copy_button=True, 
                bubble_full_width=True, 
                avatar_images=(None, "⚖️"),
                elem_classes="chatbot-container"
            )
            
            with gr.Row():
                user_msg = gr.Textbox(
                    show_label=False, 
                    placeholder="I need a Rent Agreement for Mumbai...", 
                    scale=9, 
                    container=False,
                    autofocus=True
                )
                submit_btn = gr.Button("Draft ➤", variant="primary", scale=1)
                stop_btn = gr.Button("Stop 🛑", variant="stop", scale=1, visible=False)

            with gr.Accordion("Advanced Settings", open=False):
                with gr.Row():
                    temp_s = gr.Slider(0.01, 1.0, 0.05, step=0.01, label="Temperature")
                    top_p_s = gr.Slider(0.1, 1.0, 0.9, step=0.05, label="Top P")
                    max_toks = gr.Slider(500, 4096, 2048, step=100, label="Max Tokens")

    # Wire up interactions
    submit_event = submit_btn.click(
        fn=setup_chat, inputs=[user_msg, chatbot], outputs=[user_msg, chatbot, submit_btn, stop_btn]
    ).then(
        fn=process_interaction, inputs=[chatbot, temp_s, top_p_s, max_toks], outputs=[chatbot]
    ).then(
        fn=finalize_chat, outputs=[user_msg, submit_btn, stop_btn]
    )

    user_msg.submit(
        fn=setup_chat, inputs=[user_msg, chatbot], outputs=[user_msg, chatbot, submit_btn, stop_btn]
    ).then(
        fn=process_interaction, inputs=[chatbot, temp_s, top_p_s, max_toks], outputs=[chatbot]
    ).then(
        fn=finalize_chat, outputs=[user_msg, submit_btn, stop_btn]
    )

    stop_btn.click(fn=None, cancels=[submit_event])

if __name__ == "__main__":
    print("🚀 Launching LegalBuddy Pro (Full-Screen Chat)...")
    demo.queue().launch(share=True, server_port=7865)