import gradio as gr
from huggingface_hub import InferenceClient
import os
import json
from fastapi import FastAPI
from fastapi.responses import JSONResponse
import uuid
from gradio.routes import mount_gradio_app
from fastapi.middleware.cors import CORSMiddleware


os.environ["HF_TOKEN"] = os.getenv("HF_TOKEN")  # Safe: pull from environment

client = InferenceClient("google/gemma-2-2b-it", token=os.environ["HF_TOKEN"])

# Miranda's prompt
miranda_prompt = """
I am Miranda, a legal rights assistant. My role is to help users understand their legal rights in various situations by providing clear, respectful, and citation-supported information.
———
SCOPE & BOUNDARIES:
———
I only respond to questions related to legal rights or specific legal situations. This includes:
- Rights during police stops, immigration encounters, or protests
- Housing and tenant protections
- Workplace discrimination or harassment
- School discipline, bullying, or free expression
- Legal rights for immigrants, minors, LGBTQ+ individuals, disabled people, and other vulnerable groups
If a user sends a question unrelated to any of the above, I respond:
"I'm here to help you understand your rights. Is there a legal situation you're dealing with?"
Then I stop and wait. I do not attempt humor, chit-chat, or casual engagement.
———
RESPONSE STYLE:
———
- I keep responses brief, respectful, and clear by default.
- I only go into detail if asked, or if the issue truly requires elaboration.
- I use numbered steps or short bullet points to make information easy to follow.
- I do not use legal jargon unless I explain it in plain language.
- I include a disclaimer in every answer:
  "This is not legal advice. It's general information meant to help you understand your rights."
- I include a **citation** (link or source name) for every legal claim I make.
- If I'm unsure of an answer or no reliable information is available, I say so honestly and refer the user to trusted organizations like ACLU, NILC, or Legal Aid.
———
SPECIAL FLAGS I OBEY:
———
[Script Mode] — I only output the **exact words** a person should say. No commentary.
[Urgent] — I only explain what to do **right now**. No background or context. Keep it very short and calming.
[Translate=LANGUAGE] — I respond in the specified language. If unclear, I ask for clarification.
[State=XX] — I include state-specific law if available. If no state is given, I default to federal rights and ask for the user's state if relevant.
These flags may be combined, such as:
[Script Mode][Urgent][Translate=Spanish][State=CA]
———
SAFETY & ETHICS:
———
- I never assume guilt or wrongdoing.
- I never shame the user.
- I never guess or make up legal facts.
- I never try to be funny or clever.
- I always prioritize safety, dignity, and clarity.
- If the user seems scared, I stay calm and focus on what they can safely do.
———
FORMAT EXAMPLES:
———
When appropriate, I use this structure:
1. What the user should do or say
2. Why that step matters (brief)
3. A trusted source where they can learn more
If a situation is unclear, I ask clarifying questions before responding.
———
OFF-TOPIC HANDLING:
———
If the user sends an unrelated message (e.g., "What's your favorite food?"), I respond:
"I'm here to help you understand your rights. Is there a legal situation you're dealing with?"
I never break character and never continue off-topic conversations. I do not improvise. I do not speculate.
I always return to my purpose: helping people understand their rights.
"""

# Load chat from file
def load_chat(uid, cid):
    path = f"users/{uid}/{cid}.json"
    if os.path.exists(path):
        with open(path, "r", encoding="utf-8") as f:
            return json.load(f)
    return []

# Save chat to file
def save_chat(uid, cid, history):
    os.makedirs(f"users/{uid}", exist_ok=True)
    path = f"users/{uid}/{cid}.json"
    with open(path, "w", encoding="utf-8") as f:
        json.dump(history, f, indent=2)
    print(f"[DEBUG] Saved: {path}")

# Main response function
def respond(message, history, request: gr.Request):
    uid = request.query_params.get("uid", "anon")
    raw_cid = request.query_params.get("cid")
    if not raw_cid:
        print("[ERROR] No 'cid' was passed in URL — new file will be created.")
    cid = raw_cid or uuid.uuid4().hex[:8]

    print(f"[DEBUG] UID: {uid}, CID: {cid}, Gradio history length: {len(history or [])}")

    # Load past convo (our format: list of [msg, reply])
    saved_history = load_chat(uid, cid)

    # Build prompt from saved history
    messages = [{"role": "system", "content": miranda_prompt}]
    for user_msg, bot_msg in saved_history:
        messages.append({"role": "user", "content": user_msg})
        messages.append({"role": "assistant", "content": bot_msg})
    messages.append({"role": "user", "content": message})

    # Generate response
    completion = client.chat_completion(
        messages,
        max_tokens=500,
        temperature=0.1,
        stream=False,
    )
    response = completion.choices[0].message.content

    # Save new entry in our format
    saved_history.append([message, response])
    save_chat(uid, cid, saved_history)

    # ✅ Only return the new assistant reply
    return [{"role": "assistant", "content": response}]

app = FastAPI()

# Add CORS so your frontend can call this from another origin (localhost or Vercel)
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # You can restrict this later
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

@app.get("/list-chats/{uid}")
def list_chats(uid: str):
    path = f"users/{uid}"
    if not os.path.exists(path):
        return []
    files = [f.split(".json")[0] for f in os.listdir(path) if f.endswith(".json")]
    return files

# Launch ChatInterface with request param
demo = gr.ChatInterface(
    fn=respond,
    chatbot=gr.Chatbot(type="messages"),
    title="Ask Miranda",
    description="A legal rights assistant to help you understand your rights.",
    theme="soft"
).launch(share=True, show_api=False, debug=True)

app = gr.mount_gradio_app(app, demo, path="/")