agent.py

import random
import string
import re
import datetime
import base64
import pandas as pd
from pathlib import Path
import os
import logging
from langchain_core.tools import tool
from langchain_groq import ChatGroq
from langgraph.prebuilt import create_react_agent
from langgraph.checkpoint.memory import MemorySaver

import cv2
import numpy as np
from PIL import Image

logger = logging.getLogger(__name__)

# Global variables
_SESSIONS = None
_CURRENT_THREAD_ID = None
_API_CLIENT = None
_RETRIEVER = None

def set_session_storage(sessions_dict: dict) -> None:
    global _SESSIONS
    _SESSIONS = sessions_dict

def set_current_thread_id(thread_id: str) -> None:
    global _CURRENT_THREAD_ID
    _CURRENT_THREAD_ID = thread_id

def set_api_client(api_client) -> None:
    global _API_CLIENT
    _API_CLIENT = api_client

def get_latest_image(thread_id: str) -> dict | None:
    """Get latest image from external API"""
    if _API_CLIENT is None:
        return None
    
    result = _API_CLIENT.get_patient_image(thread_id)
    if result.get('ok') and result.get('base64'):
        # Decode base64 image
        try:
            img_data = base64.b64decode(result['base64'])
            image = Image.open(io.BytesIO(img_data))
            return {"image": image}
        except Exception as e:
            logger.error(f"Error decoding image: {e}")
            return None
    return None

def _store_image_blobs(blobs: dict) -> None:
    """Store image blobs in session"""
    if not _SESSIONS or not _CURRENT_THREAD_ID:
        return
    if _CURRENT_THREAD_ID not in _SESSIONS:
        _SESSIONS[_CURRENT_THREAD_ID] = {}
    existing = _SESSIONS[_CURRENT_THREAD_ID].get("_image_blobs", {})
    existing.update(blobs)
    _SESSIONS[_CURRENT_THREAD_ID]["_image_blobs"] = existing
    names = list(_SESSIONS[_CURRENT_THREAD_ID].get("_session_image_names", []))
    for n in blobs.keys():
        if n not in names:
            names.append(n)
    _SESSIONS[_CURRENT_THREAD_ID]["_session_image_names"] = names

# System Prompt
SYSTEM_PROMPT = """You are DermaScan AI, a dermatology assistant on a dataset of 1000 skin cancer patients. Built by Eng. Youssef Bastawisy.

IMAGES: handled by tools only, on the image already uploaded this session. NEVER say you can't see/read images. On any image request, call the matching tool immediately:
- segment_my_image (segment/mask/overlay), reanalyze_my_image (analyze/re-analyze), get_my_uploaded_image (show/view photo)
No image is ever downloaded from the internet or from a patient_id — matching an uploaded image to a dataset record is handled internally by the tool. NEVER mention how the matching works internally (no "feature vectors", "global_features", "similarity", "threshold", "columns", etc.) — just state plainly whether a matching record was found or not.
If a tool says no matching record was found (new patient), you MUST say clearly this is a new patient not in the database, and STOP there for anything hereditary/genetic — never mention a mutation, family history, or risk score for them, and never reuse a different (previously matched) patient's name or data for this new image.
After a tool returns [ANALYSIS_RESULT:{data}], discuss only that data (label, confidence, affected area) — never say "shown"/"displayed".

LANGUAGE: reply in the same language/dialect the user just used (Arabic or English). Never mix in Russian/Thai/Chinese/Japanese/Korean script — if unsure of a term, describe simply instead.

PATIENT MODE STYLE: no jargon (no "asymmetry index", "ABCDE"). Plain everyday words, max 4-6 short sentences/bullets. Never mix English medical terms into Arabic sentences.

ROLE ACCESS:
- You operate in PATIENT mode only. Never reveal other patients' data, only their own uploaded result.
- If the result is Malignant: stress urgency + next steps + suggest the "📅 Book appointment" button (create a referral ticket only if explicitly asked).
- If the result is Benign: prevention + monitoring tips.

MATCHED PATIENT RECORD: if this block appears, the uploaded image matched an existing patient (same person uploading). Use it to answer diagnosis/family-history questions directly and specifically — summarize relevant parts, don't dump verbatim.

STRICT RAG RULE: your knowledge is limited to (1) patient records via matched block, (2) knowledge-base via search_knowledge_base, (3) AI image results, (4) this conversation. Never invent symptoms, history, numbers. If info isn't available say exactly: "This information is not available in the current patient record." (or Arabic equivalent).

TOOLS:
- search_knowledge_base is the default first call for general clinical/knowledge questions (cite the source filename).
- Use segment_my_image for segmentation/mask/overlay requests.
- Use reanalyze_my_image for classification + segmentation analysis.
- STOP after one tool call for a simple lookup, and at most two tool calls for a combined request.
- Never call the same tool twice.
- Never call a second different tool just because the first one returned a valid answer.
- Keep replies SHORT by default (2–4 sentences) unless the user explicitly asks for more detail.

QUICK ROUTING (pick one, don't overthink):
- General symptom/skin question, no specific image → search_knowledge_base only.
- "my/this image", "analyze it", "the result" → segment_my_image/reanalyze_my_image/get_my_uploaded_image.
- Own history/family/risk → answer from the [MATCHED PATIENT RECORD] block if present; if absent, say the info isn't available yet.

Booking is via the "📅 Book appointment" button. Never print booking IDs, tickets, or debug/system text yourself.

ANALYSIS MARKER: if a tool response starts with `[ANALYSIS_RESULT:{...}]`, keep that exact marker at the very start of your reply, unmodified. After it, briefly cover: classification + confidence, affected area, what it means in plain terms, relevant family/hereditary context if present, and next steps. Keep it short — a routine benign result needs only 2-3 sentences; a malignant one with family history can go a bit longer. Remind patients once (not every message) that this doesn't replace a professional diagnosis.

If a user gives a local file path, tell them to use the 📎 upload button instead.
"""

# Tools
@tool
def search_knowledge_base(query: str) -> str:
    """Search the dermatology knowledge base for clinical information."""
    if _RETRIEVER is None:
        return "Retriever not initialized."
    docs = _RETRIEVER.invoke(query)
    if not docs:
        return "No relevant information found."
    formatted = []
    for i, d in enumerate(docs, 1):
        source = d.metadata.get("source", "unknown").replace("\\", "/").split("/")[-1]
        formatted.append(f"[Source {i}: {source}]\n{d.page_content}")
    return "\n\n---\n\n".join(formatted)

@tool
def get_my_uploaded_image() -> str:
    """[PATIENT MODE] Tell patient to use the Show My Image button."""
    return "Your image is stored for this session. Use the 📷 Show My Image button to view it."

@tool
def segment_my_image() -> str:
    """[PATIENT MODE] Run U-Net segmentation on the patient's uploaded image."""
    if not _SESSIONS or not _CURRENT_THREAD_ID:
        return "No session found."
    
    if _API_CLIENT is None:
        return "API client not initialized."
    
    try:
        result = _API_CLIENT.segment_patient_image(_CURRENT_THREAD_ID)
        
        if not result.get('ok'):
            return f"Error: {result.get('error', 'Segmentation failed')}"
        
        # Store images
        if 'images' in result:
            _store_image_blobs(result['images'])
        
        import json as _json
        result_data = {
            "label": "",
            "confidence_pct": 0,
            "infection_pct": result.get('infection_pct', 0)
        }
        marker = f"[ANALYSIS_RESULT:{_json.dumps(result_data)}]"
        
        return f"{marker}\n\nSEGMENTATION OF YOUR IMAGE\nAffected Area: {result['infection_pct']}%"
        
    except Exception as e:
        logger.error(f"Error in segment_my_image: {e}", exc_info=True)
        return f"Error segmenting image: {str(e)}"

@tool
def reanalyze_my_image() -> str:
    """[PATIENT MODE] Full AI analysis (classification + segmentation) on the patient's
    most recently uploaded image."""
    if not _SESSIONS or not _CURRENT_THREAD_ID:
        return "No session found."
    
    if _API_CLIENT is None:
        return "API client not initialized."
    
    try:
        result = _API_CLIENT.reanalyze_patient_image(_CURRENT_THREAD_ID)
        
        if not result.get('ok'):
            return f"Error: {result.get('error', 'Analysis failed')}"
        
        # Store images
        if 'images' in result:
            _store_image_blobs(result['images'])
        
        import json as _json
        result_data = {
            "label": result.get("label", ""),
            "confidence_pct": result.get("confidence_pct", 0),
            "infection_pct": result.get("infection_pct", 0),
        }
        marker = f"[ANALYSIS_RESULT:{_json.dumps(result_data)}]"
        
        lines = [
            marker,
            "",
            "ANALYSIS OF YOUR IMAGE",
            f"Classification: {result['label']} ({result['confidence_pct']}%)",
            f"Affected Area: {result['infection_pct']}%"
        ]
        
        return "\n".join(lines)
        
    except Exception as e:
        logger.error(f"Error in reanalyze_my_image: {e}", exc_info=True)
        return f"Error analyzing image: {str(e)}"

@tool
def create_referral_ticket(issue_summary: str, urgency: str = "routine") -> str:
    """Create a dermatologist referral ticket."""
    ticket_id = "DS-" + "".join(random.choices(string.digits, k=6))
    times = {
        "routine":   "within 2–4 weeks",
        "urgent":    "within 48–72 hours",
        "emergency": "within 24 hours — go to nearest dermatology clinic",
    }
    timeframe = times.get(urgency, "within 2–4 weeks")
    
    DEFAULT_CLINIC = {
        "name": "مركز القاهرة للأمراض الجلدية",
        "address": "شارع التحرير، الدقي، الجيزة",
        "phone": "01011112222",
    }
    
    return (
        f"Referral created. Ticket: {ticket_id} | Urgency: {urgency.upper()} | "
        f"Timeframe: {timeframe} | Summary: {issue_summary} | "
        f"Clinic: {DEFAULT_CLINIC['name']} — {DEFAULT_CLINIC['address']}"
    )

def build_agent(retriever, model_name: str = "llama-3.1-8b-instant", temperature: float = 0.25):
    global _RETRIEVER
    _RETRIEVER = retriever

    groq_key = os.environ.get("GROQ_API_KEY")
    if not groq_key:
        raise ValueError("GROQ_API_KEY not set in environment variables.")

    llm = ChatGroq(
        groq_api_key=groq_key,
        model_name=model_name,
        temperature=temperature,
        max_tokens=512,
        request_timeout=120,
    )

    tools = [
        search_knowledge_base,
        get_my_uploaded_image,
        segment_my_image,
        reanalyze_my_image,
        create_referral_ticket,
    ]
    
    memory = MemorySaver()

    return create_react_agent(
        model=llm,
        tools=tools,
        prompt=SYSTEM_PROMPT,
        checkpointer=memory,
    )