
import os
import time
import json
import re
from typing import Any, Dict, List, Tuple

import streamlit as st
from PIL import Image
import pytesseract
import PyPDF2
import requests


st.set_page_config(page_title="AI Prescription Verifier", page_icon="💊", layout="wide")
st.markdown("""
<style>
html, body, [class*="css"] {font-family: Inter, system-ui, -apple-system, Segoe UI, Roboto, sans-serif;}
h1,h2,h3 {letter-spacing:.2px}
.hero{background:linear-gradient(135deg,#0f2027,#203a43,#2c5364);color:#e6f7ff;border-radius:16px;padding:18px 20px;margin-bottom:14px}
.badge{display:inline-block;font-size:12px;border:1px solid rgba(0,230,230,.45);color:#aef;border-radius:999px;padding:2px 8px;margin-left:8px}
.card{background:rgba(255,255,255,.04);border:1px solid rgba(255,255,255,.08);border-radius:14px;padding:14px}
.bot{background:rgba(0,230,230,.08);border-left:4px solid #00e6e6;padding:10px;border-radius:10px}
.small{opacity:.8;font-size:12px}
textarea, .stTextArea textarea{background:rgba(255,255,255,.06)!important;color:#fff!important;border:1px solid rgba(255,255,255,.12)!important}
</style>
""", unsafe_allow_html=True)


with st.sidebar:
    st.header("🔑 API & Models")
    st.caption("Keys are kept in this session only.")
    hf_token = st.text_input("Hugging Face API Key", type="password", value=os.getenv("HF_TOKEN",""))
    st.divider()

    st.subheader("🤖 Hugging Face models")
    GEN_MODEL = st.selectbox("Text generation (explanation)", [
        "google/flan-t5-large", "google/flan-t5-base", "google/flan-t5-xl"
    ], index=0)

    NER_MODEL = st.selectbox("NER (entity extraction)", [
        "dslim/bert-base-NER", "kamalkraj/BioBERT-NER"
    ], index=0)

    ZS_MODEL = st.selectbox("Zero-shot verification", [
        "facebook/bart-large-mnli"
    ], index=0)

    st.divider()
    st.subheader("🧠 IBM Granite (optional)")
    use_ibm = st.toggle("Enable IBM Granite 3B chat / verify", value=False)
    ibm_key = ibm_url = None
    if use_ibm:
        ibm_key = st.text_input("IBM Bearer / API Key", type="password", value=os.getenv("IBM_API_KEY",""))
        ibm_url = st.text_input("IBM Granite chat endpoint",
                                os.getenv("IBM_GRANITE_URL","https://bam-api.res.ibm.com/v1/text/chat?model_id=ibm/granite-3b-chat-v1"))
        st.caption("If you see 401/403, check Granite access & endpoint/project settings in your IBM org.")


st.markdown("""
<div class="hero">
  <h2>💊 AI Medical Prescription Verifier <span class="badge">OCR + NER + Verification</span></h2>
  <div class="small">Upload a prescription (PDF / image / TXT). We’ll extract text, find entities, check consistency, and generate a patient-friendly explanation.</div>
</div>
""", unsafe_allow_html=True)


def _hf_headers(token:str)->Dict[str,str]:
    return {"Authorization": f"Bearer {token}"} if token else {}

def hf_call(model:str, token:str, payload:Dict[str,Any], timeout:int=60)->Any:
    try:
        url = f"https://api-inference.huggingface.co/models/{model}"
        r = requests.post(url, headers=_hf_headers(token), json=payload, timeout=timeout)
        if r.status_code != 200:
            return {"error": f"{r.status_code}: {r.text[:400]}"}
        return r.json()
    except Exception as e:
        return {"error": f"Request failed: {e}"}

def extract_text(upload) -> str:
    if not upload: return ""
    if upload.type == "text/plain":
        return upload.read().decode("utf-8", errors="ignore").strip()
    if upload.type == "application/pdf":
        txt = []
        try:
            reader = PyPDF2.PdfReader(upload)
            for p in reader.pages:
                txt.append(p.extract_text() or "")
            return "\n".join(txt).strip()
        except Exception as e:
            return f"[PDF error] {e}"
    if upload.type.startswith("image/"):
        try:
            img = Image.open(upload)
            st.image(img, caption="Uploaded image", use_column_width=True)
            return pytesseract.image_to_string(img).strip()
        except Exception as e:
            return f"[OCR error] {e}"
    return "[Unsupported file type]"

def run_ner(text:str, model:str, token:str)->List[Dict[str,Any]]:
    out = hf_call(model, token, {"inputs": text, "parameters":{"aggregation_strategy":"simple"}})
    if isinstance(out, dict) and "error" in out:
        return [{"error": out["error"]}]
    ents = out[0] if isinstance(out, list) and out and isinstance(out[0], list) else out
    simp = []
    for e in ents if isinstance(ents, list) else []:
        if not isinstance(e, dict): continue
        simp.append({
            "text": e.get("word") or e.get("entity_group") or "",
            "label": e.get("entity_group") or e.get("entity") or "",
            "score": round(float(e.get("score",0)),4),
            "start": e.get("start"),
            "end": e.get("end"),
        })
    return simp

def zero_shot_check(text:str, model:str, token:str)->Dict[str,Any]:
    
    labels = [
        "medicine dosage clear",
        "medicine frequency clear",
        "duration specified",
        "possible duplication",
        "possible interaction",
        "missing critical info",
    ]
    out = hf_call(model, token, {"inputs": text, "parameters":{"candidate_labels": labels}})
    if isinstance(out, dict) and "error" in out:
        return {"error": out["error"]}
    return out

def generate_explanation(text:str, model:str, token:str)->str:
    prompt = f"""
You are a clinical AI assistant. Read the prescription and produce a concise verification for a patient.

Prescription:
\"\"\"{text}\"\"\"

Return:
- Key entities (patient, doctor if present, medicines, dosage, frequency, duration)
- Safety checks (duplication, too-high dose, likely interaction) if any
- Simple patient explanation (2–3 lines)
- Confidence (Low/Medium/High)
"""
    out = hf_call(model, token, {"inputs": prompt, "parameters":{"max_new_tokens": 256}})
    if isinstance(out, list) and out and isinstance(out[0], dict) and "generated_text" in out[0]:
        return out[0]["generated_text"]
    if isinstance(out, dict) and "error" in out:
        return f"[Generation error] {out['error']}"
    return f"[Unexpected generation result] {str(out)[:300]}"

def ibm_chat(msg:str, url:str, key:str)->str:
    try:
        headers = {"Authorization": f"Bearer {key}", "Content-Type":"application/json"}
        payload = {"input":[{"role":"user","content":msg}]}
        r = requests.post(url, headers=headers, json=payload, timeout=60)
        if r.status_code != 200:
            return f"[IBM error] {r.status_code}: {r.text[:400]}"
        data = r.json()
        
        if isinstance(data, dict):
            if "results" in data and isinstance(data["results"], list) and data["results"]:
                return data["results"][0].get("generated_text","[No text]")
            # some endpoints use 'output' or 'generated_text'
            if "generated_text" in data: return data["generated_text"]
        return f"[IBM unexpected] {str(data)[:300]}"
    except Exception as e:
        return f"[IBM request failed] {e}"

def guess_drug_tokens(entities:List[Dict[str,Any]])->List[str]:
    cand = []
    for e in entities:
        lbl = (e.get("label") or "").upper()
        w = e.get("text") or ""
        if any(k in lbl for k in ["DRUG","CHEM","TREAT","MED"]):
            cand.append(w)
        if re.match(r"^[A-Z][a-zA-Z\-]{3,}$", w):
            cand.append(w)
    seen=set(); out=[]
    for c in cand:
        if c not in seen:
            out.append(c); seen.add(c)
    return out[:20]


left, right = st.columns([1.05, 1])
with left:
    st.markdown("### 📤 Upload or Paste Prescription")
    upload = st.file_uploader("Upload (PDF, PNG, JPG, TXT)", type=["pdf","png","jpg","jpeg","txt"])
    pasted = st.text_area("Or paste text here", height=160, placeholder="Patient: ...\nRx:\n- Paracetamol 500mg, 1-0-1 after food, 5 days\n- Amoxicillin 250mg, 1-1-1 before food, 7 days")
    extract_btn = st.button("🧾 Extract Text")
    analyze_btn = st.button("🔍 Analyze with AI")

with right:
    st.markdown("### 📄 Extracted Text")
    if "extracted" not in st.session_state: st.session_state.extracted = ""

    if extract_btn:
        if upload is None and not pasted.strip():
            st.warning("Please upload a file or paste text.")
        else:
            st.session_state.extracted = extract_text(upload) if upload else pasted.strip()

    st.text_area("Output", value=st.session_state.extracted, height=220, key="out_view")

    st.markdown("### 🧠 AI Results")
    text_for_ai = st.session_state.extracted or pasted.strip()

    if analyze_btn:
        if not hf_token:
            st.error("Add your Hugging Face API key in the sidebar.")
        elif not text_for_ai:
            st.warning("No text to analyze yet.")
        else:
            
            with st.spinner("Running NER…"):
                ents = run_ner(text_for_ai, NER_MODEL, hf_token)
            if ents and isinstance(ents, list) and isinstance(ents[0], dict) and "error" in ents[0]:
                st.error(f"NER error: {ents[0]['error']}")
            else:
                st.markdown("**Named Entities:**")
                st.dataframe(ents, use_container_width=True)
                meds = guess_drug_tokens(ents)
                if meds: st.info("Detected medicine-like terms: " + ", ".join(meds))

            
            with st.spinner("Running zero-shot checks…"):
                zs = zero_shot_check(text_for_ai, ZS_MODEL, hf_token)
            if isinstance(zs, dict) and "error" in zs:
                st.error(f"Zero-shot error: {zs['error']}")
            else:
                st.markdown("**Verification scores (zero-shot):**")
                st.json(zs)

            
            with st.spinner("Generating explanation…"):
                expl = generate_explanation(text_for_ai, GEN_MODEL, hf_token)
            st.markdown("**AI Verification (structured):**")
            st.markdown(f"<div class='bot'>{expl}</div>", unsafe_allow_html=True)


st.markdown("---")
st.subheader("💬 IBM Granite 3B Chat (optional)")
if use_ibm and ibm_key and ibm_url:
    if "chat" not in st.session_state: st.session_state.chat = []
    q = st.text_input("Ask Granite about this prescription or general meds:")
    if st.button("Send to Granite"):
        st.session_state.chat.append(("you", q))
        with st.spinner("Granite is thinking…"):
            a = ibm_chat(q if q else (text_for_ai or "Summarize the prescription."), ibm_url, ibm_key)
        st.session_state.chat.append(("granite", a))
    for role, msg in st.session_state.chat[-10:]:
        if role=="you":
            st.markdown(f"**You:** {msg}")
        else:
            st.markdown(f"<div class='bot'><b>Granite:</b> {msg}</div>", unsafe_allow_html=True)
else:
    st.caption("Enable IBM Granite in the sidebar to chat (optional).")

st.markdown("---")
st.caption("Built by Aditya • Streamlit + Tesseract OCR + Hugging Face + (optional) IBM Granite 3B")
