# -----------------------------------------------------
# 1) MUST be the FIRST command
# -----------------------------------------------------
import streamlit as st
st.set_page_config(page_title="ClauseWise", layout="wide")

# -----------------------------------------------------
# 2) All other imports AFTER set_page_config
# -----------------------------------------------------
from pypdf import PdfReader
from docx import Document
import re
from multilingual import UI_TEXT, translate_text
from backup_features import (
    extract_entities,
    extract_clauses,
    get_risks,
    fairness_score,
    alternative_clauses
)
from transformers import T5Tokenizer, T5ForConditionalGeneration

# -----------------------------------------------------
# 3) Load model
# -----------------------------------------------------
@st.cache_resource
def load_chat_model():
    tokenizer = T5Tokenizer.from_pretrained("t5-small")
    model = T5ForConditionalGeneration.from_pretrained("t5-small")
    return tokenizer, model

tokenizer, chat_model = load_chat_model()

# -----------------------------------------------------
# 4) Chat response function
# -----------------------------------------------------
def chat_response(user_msg, lang):
    msg = user_msg.lower()

    # Store conversation
    if "history" not in st.session_state:
        st.session_state["history"] = []

    # Basic rule-based NDA responses
    if "nda" in msg:
        reply = "This NDA helps protect confidential information between parties."
    elif "duration" in msg or "how long" in msg:
        reply = "NDAs commonly last 1–5 years depending on the agreement."
    elif "penalty" in msg or "violate" in msg:
        reply = "Penalties depend on governing law and agreement provisions."
    elif "can i work" in msg or "competition" in msg:
        reply = "NDAs do NOT restrict employment. That is a Non-Compete Agreement."
    elif "termination" in msg:
        reply = "Termination clauses define when the NDA ends and obligations after that."
    else:
        reply = "I can help with NDA clauses, risks, obligations, confidentiality."

    reply += "\n\n⚠️ Educational only, not legal advice."

    return translate_text(reply, lang)

# -------------------------------
# Utility Functions
# -------------------------------
def is_nda(text):
    return "nda" in text.lower() or "non-disclosure" in text.lower()

def read_pdf(file):
    reader = PdfReader(file)
    return "\n".join(page.extract_text() for page in reader.pages)

def read_docx(file):
    doc = Document(file)
    return "\n".join(p.text for p in doc.paragraphs)

def read_txt(file):
    return file.read().decode("utf-8")

# -------------------------------
# Streamlit UI
# -------------------------------
st.set_page_config(page_title="ClauseWise", layout="wide")

# Language Selector
lang = st.sidebar.selectbox("🌐 Language", ["en", "hi", "ta", "te", "kn"])

# Title
st.title("📄 ClauseWise – Multilingual NDA Analyzer")

uploaded = st.file_uploader(UI_TEXT["upload_title"][lang], type=["pdf", "txt", "docx"])

if uploaded:
    ext = uploaded.name.split(".")[-1]

    if ext == "pdf":
        text = read_pdf(uploaded)
    elif ext == "txt":
        text = read_txt(uploaded)
    elif ext == "docx":
        text = read_docx(uploaded)
    else:
        st.error(UI_TEXT["error_not_nda"][lang])
        st.stop()

    if not is_nda(text):
        st.error(UI_TEXT["error_not_nda"][lang])
        st.stop()

    st.success(UI_TEXT["success_nda"][lang])

    tabs = st.tabs([
        UI_TEXT["tab_clauses"][lang],
        UI_TEXT["tab_risks"][lang],
        UI_TEXT["tab_fairness"][lang],
        UI_TEXT["tab_entities"][lang],
        UI_TEXT["tab_alternatives"][lang],
        UI_TEXT["tab_chat"][lang]
    ])

    # -----------------------------
# TAB 1 — CLAUSE SIMPLIFICATION
# -----------------------------
with tabs[0]:
    st.header(UI_TEXT["clause_simplify"][lang])

    clauses = extract_clauses(text)

    mode = st.radio(UI_TEXT["choose_mode"][lang],
                    [UI_TEXT["eli5"][lang],
                     UI_TEXT["simple"][lang],
                     UI_TEXT["pro"][lang]])

    for c in clauses:
        st.subheader(c["title"])
        if mode == UI_TEXT["eli5"][lang]:
            st.write(translate_text(c["eli5"], lang))
        elif mode == UI_TEXT["simple"][lang]:
            st.write(translate_text(c["simple"], lang))
        else:
            # ✅ FIXED SYNTAX ERROR HERE
            st.write(translate_text(c["pro"], lang))


    # -----------------------------
    # TAB 2 — RISKS
    # -----------------------------
    with tabs[1]:
        st.header(UI_TEXT["risk_title"][lang])
        risks = get_risks(text)
        for r in risks:
            st.error("⚠️ " + translate_text(r, lang))

    # -----------------------------
    # TAB 3 — FAIRNESS
    # -----------------------------
    with tabs[2]:
        st.header(UI_TEXT["fairness_title"][lang])
        score = fairness_score(text)

        st.write(f"**{UI_TEXT['your_position'][lang]}:** {score['user']}%")
        st.write(f"**{UI_TEXT['company_position'][lang]}:** {score['company']}%")

        st.progress(score["user"] / 100)

    # -----------------------------
    # TAB 4 — ENTITIES
    # -----------------------------
    with tabs[3]:
        st.header(UI_TEXT["entities_title"][lang])
        ents = extract_entities(text)

        st.write("**Parties:**", ents["parties"])
        st.write("**Dates:**", ents["dates"])
        st.write("**Amounts:**", ents["money"])

    # -----------------------------
    # TAB 5 — ALTERNATIVES
    # -----------------------------
    with tabs[4]:
        st.header(UI_TEXT["alt_title"][lang])
        alts = alternative_clauses(text)
        for a in alts:
            st.info(translate_text(a, lang))

    # -----------------------------
    # TAB 6 — LEGAL CHAT
    # -----------------------------
    with tabs[5]:
        st.header(UI_TEXT["chat_title"][lang])

        if "chat_history" not in st.session_state:
            st.session_state.chat_history = []

        for msg in st.session_state.chat_history:
            st.chat_message(msg["role"]).write(msg["text"])

        user_msg = st.chat_input(UI_TEXT["chat_placeholder"][lang])

        if user_msg:
            st.session_state.chat_history.append({"role": "user", "text": user_msg})
            st.chat_message("user").write(user_msg)

            bot_reply = chat_response(user_msg, lang)
            st.session_state.chat_history.append({"role": "assistant", "text": bot_reply})

            st.chat_message("assistant").write(bot_reply)

# Footer Disclaimer
st.info("⚠️ ClauseWise provides **educational legal insights only** — this is **NOT legal advice**.")