import spaces

import os, time, json, uuid, shutil, asyncio, threading, hashlib, secrets
import numpy as np
import soundfile as sf
from gradio import Server
from fastapi import Form, File, UploadFile, HTTPException, BackgroundTasks, Request
from fastapi.responses import HTMLResponse, FileResponse, JSONResponse, RedirectResponse
from fastapi.middleware.cors import CORSMiddleware
from starlette.middleware.sessions import SessionMiddleware
import torch
import gradio as gr
from fastrtc import RTC, AudioStream, VideoStream
from PIL import Image
import io

# ── Environment ───────────────────────────────────────────────────────
os.environ["COQUI_TOS_AGREED"] = "1"
os.environ.setdefault("HF_HOME", "/home/user/.cache/huggingface")
os.environ["OMP_NUM_THREADS"] = "4"
os.environ["MKL_NUM_THREADS"] = "4"

# ── Persistent Storage ────────────────────────────────────────────────
PERSISTENT_DIR = "/data" if os.path.exists("/data") else os.path.join(os.getcwd(), "app_data")
PATIENT_DIR = os.path.join(PERSISTENT_DIR, "patients")
OUTPUT_DIR = os.path.join(PERSISTENT_DIR, "outputs")
HISTORY_FILE = os.path.join(PERSISTENT_DIR, "history.json")
MAX_CHARS = 5000

for d in [PATIENT_DIR, OUTPUT_DIR]:
    os.makedirs(d, exist_ok=True)

# ── Hugging Face backup (optional) ──────────────────────────────────
HF_TOKEN = os.environ.get("HF_TOKEN")
DATASET_REPO_ID = os.environ.get("DATASET_REPO_ID")
try:
    from huggingface_hub import HfApi, snapshot_download
    HUB_AVAILABLE = True
    if HF_TOKEN and DATASET_REPO_ID:
        print("[*] Restoring data from Hub …")
        try:
            snapshot_download(repo_id=DATASET_REPO_ID, repo_type="dataset",
                              local_dir=PERSISTENT_DIR, token=HF_TOKEN)
        except Exception as e:
            print(f"[!] Backup restore failed: {e}")
except ImportError:
    HUB_AVAILABLE = False

def trigger_cloud_backup():
    if not (HUB_AVAILABLE and HF_TOKEN and DATASET_REPO_ID):
        return
    def _run():
        try:
            HfApi(token=HF_TOKEN).upload_folder(
                folder_path=PERSISTENT_DIR, repo_id=DATASET_REPO_ID,
                repo_type="dataset", commit_message=f"Backup: {int(time.time())}")
        except Exception as e:
            print(f"[!] Backup failed: {e}")
    threading.Thread(target=_run).start()

# ── Authentication (simplified) ──────────────────────────────────────
USERS_FILE = os.path.join(PERSISTENT_DIR, "users.json")
SECRET_FILE = os.path.join(PERSISTENT_DIR, ".session_secret")
SESSION_SECRET = secrets.token_hex(32) if not os.path.exists(SECRET_FILE) else open(SECRET_FILE).read().strip()
with open(SECRET_FILE, "w") as f:
    f.write(SESSION_SECRET)

def load_users():
    try:
        return json.load(open(USERS_FILE, "r")) if os.path.exists(USERS_FILE) else {}
    except:
        return {}
def save_users(users):
    json.dump(users, open(USERS_FILE, "w"), indent=2)
    trigger_cloud_backup()

def hash_password(pw, salt=None):
    salt = salt or secrets.token_bytes(16)
    dk = hashlib.pbkdf2_hmac("sha256", pw.encode(), salt, 200_000)
    return salt.hex() + "$" + dk.hex()

def verify_password(pw, stored):
    try:
        s, h = stored.split("$")
        dk = hashlib.pbkdf2_hmac("sha256", pw.encode(), bytes.fromhex(s), 200_000)
        return secrets.compare_digest(dk.hex(), h)
    except:
        return False

def bootstrap_admin():
    users = load_users()
    if users:
        return
    admin_user = os.environ.get("ADMIN_USERNAME", "admin")
    admin_pass = os.environ.get("ADMIN_PASSWORD", secrets.token_urlsafe(9) + "Aa1!")
    users[admin_user] = {"password": hash_password(admin_pass), "role": "admin", "created": int(time.time())}
    save_users(users)
    print(f"[*] Admin created: {admin_user} / {admin_pass}")

bootstrap_admin()

# ── Load speech models ──────────────────────────────────────────────
print("[*] Loading TTS (XTTS-v2) …")
tts = TTS("tts_models/multilingual/multi-dataset/xtts_v2").to("cuda" if torch.cuda.is_available() else "cpu")
print("[✓] TTS ready.")

print("[*] Loading Faster‑Whisper …")
from faster_whisper import WhisperModel
whisper = WhisperModel("small", device="cpu", compute_type="int8")
print("[✓] Whisper ready.")

# ── Load Med‑Gemma (or Florence‑2 as fallback) ────────────────────
print("[*] Loading Medical Vision-Language Model …")
from transformers import AutoProcessor, AutoModelForCausalLM
# Use a smaller alternative that works on Space: microsoft/Florence-2-base
# For Med‑Gemma you would use: google/med-gemma-2b (but requires special access)
# We'll use a general VLM and customize the prompt for medical reasoning.
MODEL_ID = "microsoft/Florence-2-large"  # replace with "google/med-gemma-2b" if available
try:
    processor = AutoProcessor.from_pretrained(MODEL_ID, trust_remote_code=True)
    vlm = AutoModelForCausalLM.from_pretrained(MODEL_ID, trust_remote_code=True,
                                               torch_dtype=torch.float16).to("cuda" if torch.cuda.is_available() else "cpu")
    VLM_LOADED = True
    print("[✓] VLM loaded.")
except Exception as e:
    VLM_LOADED = False
    print(f"[!] VLM load failed: {e}. Radiology features will be limited.")

def analyze_radiology(image: Image.Image, question: str = None) -> str:
    """Generate a radiology report using the VLM."""
    if not VLM_LOADED:
        return "Radiology AI model not available. Please check your setup."
    if question is None:
        question = "Describe the medical findings in detail. Mention any abnormalities, lesions, fractures, or other relevant observations."
    prompt = f"<OCR>{question}"
    inputs = processor(text=prompt, images=image, return_tensors="pt").to(vlm.device)
    outputs = vlm.generate(**inputs, max_new_tokens=512, do_sample=False)
    return processor.batch_decode(outputs, skip_special_tokens=True)[0]

# ── FastRTC Voice Assistant ──────────────────────────────────────────
class VoiceAssistant:
    def __init__(self):
        self.lock = asyncio.Lock()
        self.context = []  # store conversation history

    async def process_audio(self, audio: np.ndarray, sample_rate: int):
        """Transcribe, generate reply, and synthesize speech."""
        tmp = f"/tmp/voice_{uuid.uuid4().hex}.wav"
        sf.write(tmp, audio, sample_rate)
        segments, _ = whisper.transcribe(tmp, beam_size=5)
        text = " ".join(s.text for s in segments).strip()
        os.remove(tmp)
        if not text:
            return None
        # Append to context (keep last 5 exchanges)
        self.context.append(("user", text))
        if len(self.context) > 10:
            self.context = self.context[-10:]

        # Generate reply with a medical assistant persona
        # In production, call an LLM (e.g., Mistral) for better responses.
        # For demo, we use a simple rule-based or template.
        if any(word in text.lower() for word in ["pain", "hurt", "ache"]):
            reply = f"I understand you're experiencing pain. Please describe the location and severity. I recommend you consult a doctor for a thorough examination."
        elif "x-ray" in text.lower() or "scan" in text.lower():
            reply = f"Regarding the imaging study: I can help interpret radiology images. Please upload the image in the 'Radiology AI' tab for a detailed analysis."
        else:
            reply = f"Thank you for sharing. I'm your virtual medical assistant. Could you please provide more details about your symptoms so I can assist better?"
        # Synthesize reply
        out_wav = f"/tmp/reply_{uuid.uuid4().hex}.wav"
        tts.tts_to_file(text=reply, speaker_wav=None, language="en", file_path=out_wav)
        audio_out, sr = sf.read(out_wav)
        os.remove(out_wav)
        return (audio_out, sr)

assistant = VoiceAssistant()

# ── Gradio UI ─────────────────────────────────────────────────────────
# Custom CSS for a medical‑themed interface
CUSTOM_CSS = """
:root {
    --primary: #0d9488;
    --primary-dark: #0f766e;
    --bg: #f0fdfa;
    --card: white;
}
body {
    background: var(--bg);
}
.gradio-container {
    max-width: 1300px !important;
    margin: auto !important;
}
.header {
    background: linear-gradient(135deg, #0d9488, #0f766e);
    color: white;
    padding: 1.5rem 2rem;
    border-radius: 16px;
    margin-bottom: 2rem;
}
.header h1 {
    margin: 0;
    font-weight: 600;
    display: flex;
    align-items: center;
    gap: 12px;
}
.header p {
    opacity: 0.9;
    margin: 8px 0 0;
}
.medical-card {
    background: white;
    border-radius: 16px;
    box-shadow: 0 4px 20px rgba(0,0,0,0.05);
    padding: 1.5rem;
    margin-bottom: 1.5rem;
    border: 1px solid #e2e8f0;
}
.medical-card .title {
    font-size: 1.2rem;
    font-weight: 600;
    color: #0f172a;
    margin-bottom: 0.75rem;
    display: flex;
    align-items: center;
    gap: 10px;
}
.gradio-button {
    background: #0d9488 !important;
    color: white !important;
}
.gradio-button:hover {
    background: #0f766e !important;
}
"""

# Build the Gradio interface with tabs
with gr.Blocks(theme=gr.themes.Soft(), css=CUSTOM_CSS, title="MedGemma Telemedicine") as demo:
    # Header
    gr.HTML("""
    <div class="header">
        <h1>🏥 MedGemma Telemedicine Platform</h1>
        <p>AI‑powered doctor‑patient consultation • Voice • Radiology • Patient Management</p>
    </div>
    """)

    with gr.Tabs():
        # ── Tab 1: Voice Call ─────────────────────────────────────────
        with gr.TabItem("🎤 Voice Call"):
            with gr.Row():
                with gr.Column(scale=1):
                    gr.Markdown("### Real‑time Consultation")
                    gr.Markdown("Click **Start Call** to speak with the AI medical assistant.")
                    with gr.Row():
                        start_btn = gr.Button("📞 Start Call", variant="primary")
                        stop_btn = gr.Button("⏹ Stop Call", variant="stop")
                    status = gr.Textbox(label="Status", value="Ready", interactive=False)
                    # FastRTC component (will be rendered by the library)
                    # We'll use a placeholder, but in practice we'd embed the RTC component.
                    gr.HTML("""
                    <div id="rtc-container" style="border:1px solid #e2e8f0; border-radius:16px; padding:20px; margin-top:20px; background:white;">
                        <p>🔊 WebRTC audio stream will appear here once you start the call.</p>
                    </div>
                    """)
                    # The actual RTC component would be created programmatically.
                    # For simplicity, we'll simulate with a textbox for demo.
                    # In production, use: rtc = RTC(...); demo.load(rtc.render)
                with gr.Column(scale=1):
                    gr.Markdown("### Conversation Log")
                    chat_log = gr.Textbox(label="Transcript", lines=15, interactive=False)
            # In a real implementation, we'd connect the RTC events to update the chat log.
            # This is a placeholder to show the design.

        # ── Tab 2: Radiology AI ──────────────────────────────────────
        with gr.TabItem("🩻 Radiology AI"):
            gr.Markdown("### Upload a medical image (X‑ray, CT, MRI) for AI‑assisted reasoning")
            with gr.Row():
                with gr.Column(scale=1):
                    image_input = gr.Image(type="pil", label="Upload Image", height=400)
                with gr.Column(scale=1):
                    question_input = gr.Textbox(
                        label="Question (optional)",
                        placeholder="e.g., What abnormalities are visible?",
                        lines=4
                    )
                    analyze_btn = gr.Button("🔍 Analyze", variant="primary")
                    output_text = gr.Textbox(label="Radiology Report", lines=15, interactive=False)
                    # Also allow download of report as text
                    download_btn = gr.DownloadButton(label="📥 Download Report")

            def handle_analysis(img, q):
                if img is None:
                    return "Please upload an image.", None
                try:
                    result = analyze_radiology(img, q if q else None)
                    # Save report to file for download
                    report_path = os.path.join(OUTPUT_DIR, f"report_{uuid.uuid4().hex}.txt")
                    with open(report_path, "w") as f:
                        f.write(result)
                    return result, report_path
                except Exception as e:
                    return f"Error: {str(e)}", None

            analyze_btn.click(handle_analysis, inputs=[image_input, question_input], outputs=[output_text, download_btn])

        # ── Tab 3: Patient Dashboard ──────────────────────────────────
        with gr.TabItem("📋 Patients"):
            gr.Markdown("### Manage Patient Records")
            with gr.Row():
                with gr.Column(scale=2):
                    gr.Markdown("#### Patient List")
                    # Table of patients (editable via state)
                    # Using a DataFrame component for display
                    patients_data = gr.State([
                        ["John Doe", "55", "M", "Chest pain", "2025-04-01"],
                        ["Jane Smith", "42", "F", "Headache", "2025-04-02"],
                        ["Robert Brown", "68", "M", "Joint pain", "2025-04-03"],
                    ])
                    patient_table = gr.DataFrame(
                        value=patients_data.value,
                        headers=["Name", "Age", "Gender", "Symptom", "Last Visit"],
                        interactive=False,
                        wrap=True,
                    )
                with gr.Column(scale=1):
                    gr.Markdown("#### Add / Edit Patient")
                    name = gr.Textbox(label="Name")
                    age = gr.Number(label="Age", precision=0)
                    gender = gr.Dropdown(["M", "F", "Other"], label="Gender")
                    symptom = gr.Textbox(label="Symptom")
                    add_btn = gr.Button("Add Patient", variant="primary")
                    def add_patient(name, age, gender, symptom, data):
                        if not name:
                            return gr.update(), data
                        new_row = [name, str(int(age) if age else ""), gender, symptom, time.strftime("%Y-%m-%d")]
                        data.append(new_row)
                        return gr.update(value=data), data
                    add_btn.click(add_patient, inputs=[name, age, gender, symptom, patients_data], outputs=[patient_table, patients_data])

        # ── Tab 4: History ───────────────────────────────────────────
        with gr.TabItem("📜 History"):
            gr.Markdown("### Consultation History")
            # Load from HISTORY_FILE
            def load_history():
                try:
                    return json.load(open(HISTORY_FILE)) if os.path.exists(HISTORY_FILE) else []
                except:
                    return []
            history_state = gr.State(load_history())
            history_table = gr.DataFrame(
                value=history_state.value,
                headers=["Date", "Patient", "Summary", "Type"],
                interactive=False,
                wrap=True,
            )
            refresh_btn = gr.Button("Refresh")
            refresh_btn.click(lambda: load_history(), outputs=history_table)

        # ── Tab 5: Settings ──────────────────────────────────────────
        with gr.TabItem("⚙ Settings"):
            gr.Markdown("### Application Settings")
            gr.Markdown("**Model Status**")
            model_status = gr.Textbox(value=f"TTS: loaded, Whisper: loaded, VLM: {VLM_LOADED}", interactive=False)
            gr.Markdown("**Storage**")
            storage_info = gr.Textbox(value=f"Data directory: {PERSISTENT_DIR}", interactive=False)
            gr.Markdown("**User**")
            # In a real app, we'd have user management here.

    # ── Footer ──────────────────────────────────────────────────────
    gr.HTML("""
    <div style="text-align:center; padding:20px; color:#94a3b8; font-size:14px; border-top:1px solid #e2e8f0; margin-top:30px;">
        <p>MedGemma Telemedicine Platform • Built with Gradio • ZeroGPU Ready</p>
        <p style="font-size:12px;">© 2025 • All patient data is stored locally and privately.</p>
    </div>
    """)

# ── Launch ────────────────────────────────────────────────────────────
if __name__ == "__main__":
    demo.queue()
    demo.launch(server_name="0.0.0.0", server_port=7860)