import gradio as gr
import base64
import re
import os
import tempfile
from openai import OpenAI
from sarvamai import SarvamAI
from sarvamai.play import save

# ─── API KEYS ─────────────────────────────────────────────────────────────────
OPENAI_KEY  = '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'
SARVAM_KEY  = 'sk_lnnt2cq6_BsU3lYlvinKhKfR4x2v8WqGC'

# ─── SYSTEM PROMPT ────────────────────────────────────────────────────────────
SYSTEM_PROMPT = """
You are a medical report summarization assistant.
STRICT RULES:
1. Only process medical report images (lab reports, prescriptions, scans, test reports).
2. If the image is NOT a medical report, respond ONLY with:
   "ERROR: The uploaded image is not a valid medical report."
3. Summarize in VERY SIMPLE language so a non-medical person can understand.
4. Avoid medical jargon. If needed, explain it simply.
Your response MUST follow this EXACT structure with these two section headers:
## DETAILED SUMMARY
Write a thorough, easy-to-understand breakdown of the report.
Cover: key values, abnormal findings (in simple terms), what each result means for
the person's health, and anything they should be aware of.
End this section with: "AI Generated Summary: Always consult a doctor before taking any further steps."
## CONCISE SUMMARY
Write a short paragraph of AT LEAST 40 WORDS and NO MORE THAN 50 WORDS.
It must cover ALL key findings — normal results, abnormal values, and what the person
should watch out for. Miss nothing important. Every aspect of the report must be reflected.
"""

# ─── HELPERS ──────────────────────────────────────────────────────────────────
def encode_image(path: str) -> str:
    with open(path, "rb") as f:
        return base64.b64encode(f.read()).decode("utf-8")

def extract_section(text: str, header: str) -> str:
    pattern = rf"##\s*{re.escape(header)}\s*\n(.*?)(?=\n##\s|\Z)"
    match = re.search(pattern, text, re.DOTALL | re.IGNORECASE)
    return match.group(1).strip() if match else ""

def generate_tts(text: str, api_key: str) -> tuple[str | None, str]:
    """Generate TTS audio using Sarvam AI and return (path_to_wav, status_message)."""
    try:
        if not text or not text.strip():
            return None, "No text provided for TTS"

        client = SarvamAI(api_subscription_key=api_key)

        audio = client.text_to_speech.convert(
            target_language_code="en-IN",
            text=text.strip(),
            model="bulbul:v3",
            speaker="varun",
        )

        tmp = tempfile.NamedTemporaryFile(
            delete=False, suffix=".wav", dir=tempfile.gettempdir()
        )
        tmp.close()

        save(audio, tmp.name)

        if not os.path.exists(tmp.name) or os.path.getsize(tmp.name) == 0:
            return None, "TTS generated empty audio"

        return tmp.name, "TTS generated successfully"

    except Exception as e:
        return None, f"[TTS ERROR] {str(e)}"

# ─── CORE LOGIC ───────────────────────────────────────────────────────────────
def analyse_report(image_path):
    if image_path is None:
        return (
            "<div class='error-box'>⚠️ PLEASE UPLOAD A MEDICAL DOCUMENT FIRST.</div>",
            "", None, ""
        )

    openai_client = OpenAI(api_key=OPENAI_KEY)
    img_b64       = encode_image(image_path)

    # ── GPT-4o call ──────────────────────────────────────────────────────────
    full_output = ""
    try:
        with openai_client.responses.stream(
            model="gpt-4o",
            input=[
                {"role": "system", "content": SYSTEM_PROMPT},
                {
                    "role": "user",
                    "content": [
                        {"type": "input_text",  "text": "Analyze and summarize this medical report"},
                        {"type": "input_image", "image_url": f"data:image/jpeg;base64,{img_b64}"}
                    ]
                }
            ]
        ) as stream:
            for event in stream:
                if event.type == "response.output_text.delta":
                    full_output += event.delta
            final_response = stream.get_final_response()
    except Exception as e:
        return (f"<div class='error-box'>❌ OPENAI ERROR: {e}</div>", "", None, "")

    if full_output.strip().startswith("ERROR:"):
        return (
            f"<div class='error-box'>⚠️ {full_output.strip()}</div>",
            "", None, ""
        )

    detailed_summary = extract_section(full_output, "DETAILED SUMMARY")
    summary_50       = extract_section(full_output, "CONCISE SUMMARY")

    words = summary_50.split()
    if len(words) > 50:
        summary_50 = " ".join(words[:50]) + "..."

    # ── Token cost ────────────────────────────────────────────────────────────
    usage      = final_response.usage
    total_cost = (usage.input_tokens * 2.5 + usage.output_tokens * 10.0) / 1_000_000
    metrics_html = f"""
    <div class='metrics-row'>
        <div class='chip'><span class='val'>{usage.input_tokens:,}</span><span class='lbl'>INPUT TOKENS</span></div>
        <div class='chip'><span class='val'>{usage.output_tokens:,}</span><span class='lbl'>OUTPUT TOKENS</span></div>
        <div class='chip'><span class='val'>${total_cost:.5f}</span><span class='lbl'>ANALYSIS COST</span></div>
    </div>"""

    # ── Sarvam AI TTS ─────────────────────────────────────────────────────────
    audio_path, tts_status = generate_tts(summary_50, SARVAM_KEY)

    safe_detailed = detailed_summary.replace("<", "&lt;").replace(">", "&gt;")
    safe_concise  = summary_50.replace("<", "&lt;").replace(">", "&gt;")

    detailed_html = f"""
    <div class='result-card detailed'>
        <div class='card-header'>
            <div class='card-icon blue'>❖</div>
            <div class='card-title'>COMPREHENSIVE BREAKDOWN</div>
        </div>
        <div class='card-body'>{safe_detailed}</div>
        {metrics_html}
    </div>"""

    # Animated Waveform integration
    waveform_html = """
    <div class="waveform-container">
        <div class="bar"></div><div class="bar"></div><div class="bar"></div><div class="bar"></div><div class="bar"></div>
    </div>
    """ if audio_path else ""

    tts_indicator = (
        f'<div class="audio-badge ready">{waveform_html} AUDIO SIGNAL ACTIVE</div>'
        if audio_path else
        '<div class="audio-badge failed">✗ AUDIO SIGNAL OFFLINE</div>'
    )

    concise_html = f"""
    <div class='result-card concise'>
        <div class='card-header' style="justify-content: space-between;">
            <div style="display:flex; align-items:center; gap:10px;">
                <div class='card-icon coral'>⚡</div>
                <div class='card-title'>QUICK BRIEFING</div>
            </div>
            {tts_indicator}
        </div>
        <div class='card-body focus-text'>{safe_concise}</div>
    </div>"""

    disclaimer = """
    <div class='disclaimer'>
        <span style="font-size:18px;">⚠️</span>
        <span><strong>MEDICAL OVERRIDE WARNING:</strong> THIS INTELLIGENCE IS AI-GENERATED FOR INFORMATIONAL PURPOSES. IT DOES NOT REPLACE PROFESSIONAL MEDICAL ADVICE. ALWAYS CONSULT YOUR HEALTHCARE PROVIDER.</span>
    </div>"""

    return detailed_html, concise_html, audio_path, disclaimer


# ─── CUSTOM CSS (Neon Cyberpunk UI/UX) ─────────────────────────────────────────
CSS = """
@import url('https://fonts.googleapis.com/css2?family=Outfit:wght@300;400;500;700&family=Share+Tech+Mono&display=swap');

:root {
    --bg-main: #000000;
    --bg-card: #080808;
    --primary: #00FFFF; /* Neon Cyan */
    --primary-glow: rgba(0, 255, 255, 0.6);
    --primary-soft: rgba(0, 255, 255, 0.1);
    
    --accent: #FF00FF; /* Neon Magenta */
    --accent-glow: rgba(255, 0, 255, 0.6);
    --accent-soft: rgba(255, 0, 255, 0.1);
    
    --border-soft: rgba(0, 255, 255, 0.3);
    
    --text-neon: #E0FFFF; /* Bright icy blue for main readability */
    --text-muted: #008B8B;
    
    --shadow-glow: 0 0 10px var(--primary-glow), inset 0 0 10px var(--primary-soft);
    --shadow-accent: 0 0 10px var(--accent-glow), inset 0 0 10px var(--accent-soft);
}

body, .gradio-container { 
    background: var(--bg-main) !important; 
    font-family: 'Share Tech Mono', monospace !important; 
    color: var(--text-neon) !important; 
}
.gradio-container { max-width: 1280px !important; margin: 0 auto !important; padding: 20px !important; }
footer { display: none !important; }

/* HERO SECTION */
.hero { text-align: left; padding: 20px 10px 40px; }
.hero-title {
    font-family: 'Outfit', sans-serif;
    font-size: clamp(32px, 4vw, 48px); font-weight: 700; 
    color: #FFF; margin: 0; letter-spacing: 2px; text-transform: uppercase;
    text-shadow: 0 0 8px var(--primary), 0 0 20px var(--primary);
}
.hero-title span { 
    color: #FFF; 
    text-shadow: 0 0 8px var(--accent), 0 0 20px var(--accent);
}
.hero-sub { 
    font-family: 'Share Tech Mono', monospace; font-size: 16px; 
    color: var(--primary); margin: 12px 0 0; line-height: 1.6; max-width: 550px; 
    text-shadow: 0 0 4px var(--primary);
}

/* PANELS (For single screen fitting) */
.left-panel { display: flex; flex-direction: column; gap: 20px; }
.right-panel { 
    display: flex; flex-direction: column; gap: 20px; 
    max-height: 85vh; overflow-y: auto; padding-right: 15px;
}
.right-panel::-webkit-scrollbar { width: 4px; }
.right-panel::-webkit-scrollbar-thumb { background: var(--primary); box-shadow: 0 0 5px var(--primary); border-radius: 10px; }

/* UPLOAD AREA */
.upload-panel {
    background: var(--bg-card); border: 1px dashed var(--primary);
    border-radius: 12px; overflow: hidden; transition: all 0.3s ease;
    box-shadow: var(--shadow-glow);
}
.upload-panel:hover { border: 1px solid var(--accent); box-shadow: var(--shadow-accent); }
.upload-panel .wrap { background: transparent !important; border: none !important; }
.upload-panel svg { color: var(--primary) !important; filter: drop-shadow(0 0 5px var(--primary)); }
.upload-panel span { color: var(--text-neon) !important; text-shadow: 0 0 5px var(--primary); }

/* ACTION BUTTON */
#analyse-btn {
    background: var(--bg-card) !important;
    border: 1px solid var(--primary) !important;
    border-radius: 8px !important; color: var(--primary) !important;
    font-family: 'Share Tech Mono', monospace !important; font-size: 18px !important;
    letter-spacing: 2px !important; padding: 18px !important;
    box-shadow: var(--shadow-glow) !important;
    text-shadow: 0 0 5px var(--primary) !important;
    transition: all 0.2s ease !important;
    text-transform: uppercase;
}
#analyse-btn:hover { 
    background: var(--primary-soft) !important;
    border-color: var(--accent) !important; 
    color: var(--accent) !important;
    box-shadow: var(--shadow-accent) !important; 
    text-shadow: 0 0 8px var(--accent) !important;
    transform: scale(1.02) !important;
}

/* RESULT CARDS */
.result-card {
    background: var(--bg-card); border: 1px solid var(--border-soft);
    border-radius: 12px; padding: 28px 32px;
    box-shadow: inset 0 0 15px rgba(0,255,255,0.05); position: relative;
}
.result-card::before {
    content: ''; position: absolute; top: 0; left: 0; width: 100%; height: 2px;
    background: linear-gradient(90deg, transparent, var(--primary), transparent);
}
.result-card.concise::before {
    background: linear-gradient(90deg, transparent, var(--accent), transparent);
}

.card-header { display: flex; align-items: center; gap: 12px; margin-bottom: 16px; border-bottom: 1px solid var(--border-soft); padding-bottom: 16px;}
.card-icon { width: 36px; height: 36px; border-radius: 8px; display: flex; align-items: center; justify-content: center; font-size: 18px; border: 1px solid; }
.card-icon.blue { background: var(--primary-soft); color: var(--primary); border-color: var(--primary); box-shadow: 0 0 8px var(--primary); text-shadow: 0 0 5px var(--primary); }
.card-icon.coral { background: var(--accent-soft); color: var(--accent); border-color: var(--accent); box-shadow: 0 0 8px var(--accent); text-shadow: 0 0 5px var(--accent); }
.card-title { font-size: 16px; font-weight: 700; letter-spacing: 0.1em; color: var(--primary); text-shadow: 0 0 5px var(--primary); }
.result-card.concise .card-title { color: var(--accent); text-shadow: 0 0 5px var(--accent); }

.card-body { font-family: 'Outfit', sans-serif; font-size: 16px; font-weight: 300; line-height: 1.8; color: var(--text-neon); white-space: pre-wrap; }
.focus-text { font-size: 18px; font-weight: 400; line-height: 1.6; color: #FFF; text-shadow: 0 0 4px var(--accent); }

/* METRICS */
.metrics-row { display: flex; gap: 12px; margin-top: 24px; }
.chip { flex: 1; background: #000; border-radius: 8px; padding: 12px 14px; text-align: center; border: 1px solid var(--border-soft); box-shadow: inset 0 0 8px rgba(0,255,255,0.1); }
.chip .val { display: block; font-size: 20px; font-weight: 700; color: var(--primary); text-shadow: 0 0 5px var(--primary); }
.chip .lbl { display: block; font-size: 12px; color: var(--text-muted); margin-top: 4px; }

/* AUDIO COMPONENT STYLING */
#audio-out { background: transparent !important; border: none !important; padding: 0 !important; margin-bottom: 10px !important; }
#audio-out audio { 
    width: 100% !important; border-radius: 8px !important; height: 50px !important; 
    border: 1px solid var(--accent) !important;
    box-shadow: 0 0 10px var(--accent-glow) !important;
    filter: invert(1) hue-rotate(180deg); /* Forces a dark look on default audio players */
}

/* WAVEFORM ANIMATION */
.audio-badge { display: flex; align-items: center; gap: 8px; font-size: 12px; font-weight: 700; padding: 6px 12px; border-radius: 4px; border: 1px solid; }
.audio-badge.ready { background: var(--accent-soft); color: var(--accent); border-color: var(--accent); box-shadow: 0 0 8px var(--accent); text-shadow: 0 0 4px var(--accent); }
.audio-badge.failed { background: rgba(255,0,0,0.1); color: #FF0000; border-color: #FF0000; box-shadow: 0 0 8px #FF0000; text-shadow: 0 0 4px #FF0000; }
.waveform-container { display: flex; align-items: center; gap: 2px; height: 14px; margin-right: 4px; }
.bar { width: 3px; background: var(--accent); border-radius: 1px; box-shadow: 0 0 5px var(--accent); animation: wave 1s ease-in-out infinite; }
.bar:nth-child(1) { height: 40%; animation-delay: 0.0s; }
.bar:nth-child(2) { height: 80%; animation-delay: 0.1s; }
.bar:nth-child(3) { height: 100%; animation-delay: 0.2s; }
.bar:nth-child(4) { height: 60%; animation-delay: 0.3s; }
.bar:nth-child(5) { height: 40%; animation-delay: 0.4s; }
@keyframes wave { 0%, 100% { transform: scaleY(0.4); } 50% { transform: scaleY(1); } }

/* DISCLAIMER */
.disclaimer { 
    background: rgba(255, 255, 0, 0.05); border-left: 4px solid #FFFF00; border-radius: 0 8px 8px 0; 
    padding: 16px 20px; font-size: 13px; color: #FFFF00; display: flex; align-items: center; gap: 12px; line-height: 1.5; 
    box-shadow: 0 0 10px rgba(255, 255, 0, 0.2); text-shadow: 0 0 4px rgba(255, 255, 0, 0.5);
}
.error-box {
    background: rgba(255,0,0,0.1); border: 1px solid #FF0000; padding: 15px; border-radius: 8px;
    color: #FF0000; text-shadow: 0 0 5px #FF0000; box-shadow: 0 0 10px rgba(255,0,0,0.3); text-align: center;
}

/* MOBILE RESPONSIVENESS */
@media (max-width: 900px) {
    .right-panel { max-height: none; overflow-y: visible; padding-right: 0; }
    .hero-title { font-size: 28px; }
    .card-header { flex-direction: column; align-items: flex-start; }
}
"""

# ─── GRADIO UI ────────────────────────────────────────────────────────────────
with gr.Blocks(css=CSS, theme=gr.themes.Base(), title="MediScan AI") as demo:

    with gr.Row():
        # LEFT COLUMN (Uploader & Controls)
        with gr.Column(scale=4, elem_classes="left-panel"):
            gr.HTML("""
            <div class="hero">
                <div class="hero-title">MEDI<span>SCAN</span> AI</div>
                <p class="hero-sub">INITIALIZE UPLOAD. AI CLINICAL ENGINE STANDING BY FOR PLAIN-LANGUAGE DECRYPTION & AUDIO SYNTHESIS.</p>
            </div>
            """)

            with gr.Column(elem_classes="upload-panel"):
                image_input = gr.Image(
                    type="filepath",
                    label="DRAG & DROP MEDICAL DATAPAD HERE",
                    show_label=True,
                    height=380,
                )

            analyse_btn = gr.Button("INITIALIZE SCAN & SYNTHESIZE AUDIO", elem_id="analyse-btn")

        # RIGHT COLUMN (Outputs - scrollable on desktop, fits on one screen)
        with gr.Column(scale=6, elem_classes="right-panel"):
            audio_out = gr.Audio(
                label="",
                type="filepath",
                autoplay=True,
                elem_id="audio-out",
                show_label=False
            )
            
            concise_out    = gr.HTML()
            detailed_out   = gr.HTML()
            disclaimer_out = gr.HTML()

    # Link button logic
    analyse_btn.click(
        fn=analyse_report,
        inputs=[image_input],
        outputs=[detailed_out, concise_out, audio_out, disclaimer_out],
    )

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