import gradio as gr
import requests
from transformers import pipeline

# ==========================================================
# MODEL: flan-t5-base (बेहतर जवाब!)
# ==========================================================
MODEL_ID = "google/flan-t5-small"  # ← t5-small पर परीक्षण के लिए
# MODEL_ID = "google/flan-t5-base"  # ← जब समस्या हल हो जाए तो इसे वापस स्विच कर सकते हैं
print(f"Loading {MODEL_ID}...")
generator = pipeline("text2text-generation", model=MODEL_ID, device=-1)

# ==========================================================
# PROMPT
# ==========================================================
def get_prompt(user_input):
    text = user_input.strip()
    if not text:
        return "Question: कोई प्रश्न नहीं।\nAnswer:"
    text = text.replace(" या ", " है या नहीं? ")
    if not text.endswith("?"):
        text += "?"
    return f"Question: {text}\nAnswer:"

# ==========================================================
# MODEL CALL
# ==========================================================
def call_model(user_input):
    try:
        result = generator(
            get_prompt(user_input),
            max_new_tokens=150,  # max_length की जगह इसे इस्तेमाल करें
            do_sample=False,
            temperature=0.3
        )
        answer = result[0]["generated_text"]
        if "Answer:" in answer:
            answer = answer.split("Answer:")[-1].strip()
        return answer if answer else "मैं अभी इस पर जवाब नहीं दे सकता।"
    except Exception as e:
        # Debugging: पूरी एरर मैसेज प्रिंट करें
        print(f"Error in call_model: {e}")
        return f"त्रुटि: {str(e)}"

# ==========================================================
# SIA LAYER (वही) - अभी इन्हें सीधे उपयोग नहीं करेंगे
# ==========================================================
def fact_check(text):
    try:
        q = text.split('.')[0].strip()[:80]
        if len(q) < 3: return False, None
        res = requests.get("https://en.wikipedia.org/w/api.php", params={
            "action": "opensearch", "search": q, "limit": 1, "format": "json"
        }, timeout=5).json()
        return bool(res[1]), res[1][0] if res[1] else None
    except requests.exceptions.Timeout:
        print("Wikipedia API call timed out in fact_check!")
        return False, None
    except requests.exceptions.RequestException as e:
        print(f"Wikipedia API request failed in fact_check: {e}")
        return False, None
    except Exception as e:
        print(f"An unexpected error occurred in fact_check: {e}")
        return False, None

def add_empathy(text):
    if any(w in text.lower() for w in ["मदद", "कृपया", "बताएं"]):
        return "मैं समझता/समझती हूँ कि यह महत्वपूर्ण है। " + text
    return "ध्यानपूर्वक: " + text

def remove_bias(text):
    return text.replace("idiot", "व्यक्ति").replace("मूर्ख", "अज्ञानी")

def calculate_ehi(raw, final):
    score = 50
    if "समझ" in final or "ध्यानपूर्वक" in final: score += 30
    if len(final) >= 20: score += 20
    return min(score, 100)

# ==========================================================
# MAIN - SIA Layer को बाईपास करने के लिए संशोधित
# ==========================================================
def run_demo(user_input):
    print(f"User input received: {user_input}") # नया
    if not user_input.strip():
        return "प्रश्न लिखें।", "", "EHI: 0", "None"
    
    print("Calling model (SIA bypass)...") # नया
    raw = call_model(user_input)
    print(f"Raw output from model: {raw}") # नया
    
    # SIA Layer को बाईपास कर दिया गया है
    sia_filtered_output = "SIA bypassed" # या आप raw आउटपुट ही दिखा सकते हैं
    ehi_score = "EHI: N/A"
    fact_anchor = "N/A (SIA bypassed)"

    # यदि आप चाहते हैं कि raw आउटपुट SIA-Filtered बॉक्स में भी दिखे:
    # sia_filtered_output = raw

    return raw, sia_filtered_output, ehi_score, fact_anchor

# ==========================================================
# UI
# ==========================================================
with gr.Blocks(title="SIA Infinity") as demo:
    gr.Markdown("# SIA Infinity — Ethical AI Demo\n**Free • Local • Hindi+English**")
    inp = gr.Textbox(label="प्रश्न लिखें", lines=3, placeholder="AI क्या है?\nभारत की राजधानी?")
    with gr.Row():
        raw_out = gr.Textbox(label="Raw Output", lines=4)
        sia_out = gr.Textbox(label="SIA-Filtered", lines=4)
    with gr.Row():
        ehi = gr.Textbox(label="EHI")
        fact = gr.Textbox(label="Fact Anchor")
    gr.Button("Run Demo", variant="primary").click(run_demo, inp, [raw_out, sia_out, ehi, fact])

demo.launch(server_name="0.0.0.0", server_port=7860, ssr_mode=True)