src/core/protocol_generator.py

"""
Human+ Protocol Generator
Menggunakan plain-text delimiter output — BUKAN JSON.

Alasan tidak pakai JSON:
- Protocol lengkap bisa 6000-10000 token per response
- JSON yang terpotong di tengah menyebabkan silent parse failure
- Semua section menjadi kosong tanpa error yang jelas
- Plain text dengan delimiter jauh lebih andal dan tidak ada risiko truncation
"""

import re
from core.ai_client import call_gemini, TEMP_PROTOCOL, MAX_TOKENS_PROTOCOL
from core.rag_engine import retrieve_context
from langchain_community.vectorstores import FAISS


# ── Section delimiters ────────────────────────────────────────────────
# Dipakai untuk split response Gemini menjadi dict per section
_DELIMITERS = {
    "summary":     "===SUMMARY===",
    "iv_drip":     "===IV_DRIP===",
    "supplements": "===SUPPLEMENTS===",
    "lifestyle":   "===LIFESTYLE===",
    "retest":      "===RETEST===",
}

# ── System identity ───────────────────────────────────────────────────
_SYSTEM_IDENTITY = """
Kamu adalah AI Longevity Specialist resmi dari Human+ Bali, didirikan oleh Benjamin White (@bennywhitethatsright).

Filosofi utama: "From Injured → Optimized"
Voice: Positif, empowering, seperti seorang coach yang peduli — bukan dokter yang menghakimi.
Bahasa: Indonesia yang mudah dipahami, natural, tidak terlalu formal.

Background Benjamin White: Pernah koma dan mengalami traumatic brain injury.
Bangkit melalui 8+ tahun natural medicine, neuroscience, dan biohacking.
Sekarang membantu orang lain melakukan hal yang sama di Bali.
"""

# ── Main prompt ───────────────────────────────────────────────────────
_PROTOCOL_PROMPT_TEMPLATE = """{system_identity}

═══════════════════════════════════════
HASIL LAB PASIEN:
{biomarker_summary}

═══════════════════════════════════════
KNOWLEDGE BASE HUMAN+:
{context}

═══════════════════════════════════════
INSTRUKSI:

Buat rekomendasi protocol hyperpersonal yang LENGKAP dan DETAIL berdasarkan hasil lab di atas.

Format output WAJIB menggunakan delimiter berikut — tulis persis seperti ini:

===SUMMARY===
Tulis 3-4 kalimat ringkasan temuan utama. Sebutkan marker spesifik yang paling kritis
beserta nilainya. Akhiri dengan motivasi positif "From Injured → Optimized".

===IV_DRIP===
Tulis rekomendasi IV Drip yang LENGKAP dan DETAIL:
- Nama protocol (The All-In atau varian custom)
- SEMUA komposisi dengan dosis spesifik
- Frekuensi dan durasi per sesi
- Penjelasan MENGAPA setiap komponen dipilih berdasarkan nilai lab pasien

===SUPPLEMENTS===
Tulis suplemen harian yang LENGKAP:
- Minimal 6-8 suplemen yang relevan dengan hasil lab
- Setiap suplemen: nama, dosis, waktu minum, dan alasan (kaitkan dengan nilai lab)
- Gunakan format list atau tabel yang jelas

===LIFESTYLE===
Tulis protocol lifestyle yang LENGKAP dan KONKRET:
1. Ice Bath / Cold Exposure (suhu, durasi, frekuensi, breathing protocol)
2. Movement & Exercise (jenis, frekuensi, durasi)
3. Sleep & Recovery (jam tidur, tips HRV, no-screen, dll)
4. Nutrisi & Diet (makanan yang dianjurkan/dihindari, khusus berdasarkan hasil lab)
5. Sunlight & Circadian Rhythm

===RETEST===
Tulis jadwal retest yang KONKRET:
- Timeline retest pertama (kapan)
- Marker prioritas yang perlu diperiksa
- Target angka spesifik yang ingin dicapai per marker
- Next step dan CTA untuk konsultasi langsung dengan Human+ Bali

PENTING:
- Selalu sebut nilai lab pasien secara spesifik (contoh: "Vitamin D Anda 22 ng/mL, target kami 60-70 ng/mL")
- Setiap rekomendasi HARUS terhubung ke hasil lab, bukan generic
- Tulis selengkap dan sedetail mungkin untuk setiap section
- Gunakan markdown (bold, bullet, tabel) agar mudah dibaca
- Akhiri dengan semangat "From Injured → Optimized"
"""


class ProtocolGeneratorError(Exception):
    pass


def generate_protocol(
    biomarkers: list[dict],
    vector_store: FAISS = None,
) -> dict:
    """
    Generate protocol dari biomarker list.
    Output: dict dengan keys summary, iv_drip, supplements, lifestyle, retest.
    """
    biomarker_summary = _format_biomarkers_for_prompt(biomarkers)

    try:
        context = retrieve_context(biomarker_summary, vector_store)
    except Exception as e:
        raise ProtocolGeneratorError(f"RAG retrieval gagal: {e}") from e

    prompt = _PROTOCOL_PROMPT_TEMPLATE.format(
        system_identity=_SYSTEM_IDENTITY,
        biomarker_summary=biomarker_summary,
        context=context,
    )

    try:
        raw_response = call_gemini(
            contents=[prompt],
            temperature=TEMP_PROTOCOL,
            max_tokens=MAX_TOKENS_PROTOCOL,  # 16000 — cukup untuk semua section
        )
    except Exception as e:
        raise ProtocolGeneratorError(f"Gemini protocol generation gagal: {e}") from e

    return _parse_delimiter_response(raw_response)


def generate_protocol_from_raw_text(
    extracted_text: str,
    vector_store: FAISS = None,
) -> dict:
    """Fallback: generate dari raw text (non-structured biomarkers)."""
    try:
        context = retrieve_context(extracted_text, vector_store)
    except Exception as e:
        raise ProtocolGeneratorError(f"RAG retrieval gagal: {e}") from e

    prompt = _PROTOCOL_PROMPT_TEMPLATE.format(
        system_identity=_SYSTEM_IDENTITY,
        biomarker_summary=extracted_text,
        context=context,
    )

    try:
        raw_response = call_gemini(
            contents=[prompt],
            temperature=TEMP_PROTOCOL,
            max_tokens=MAX_TOKENS_PROTOCOL,
        )
    except Exception as e:
        raise ProtocolGeneratorError(f"Gemini gagal: {e}") from e

    return _parse_delimiter_response(raw_response)


# ── Private helpers ──────────────────────────────────────────────────

def _parse_delimiter_response(raw: str) -> dict:
    """
    Parse plain-text response dengan delimiter menjadi dict per section.

    Strategi:
    1. Split berdasarkan ===SECTION=== delimiter
    2. Map ke dict key yang sesuai
    3. Fallback: jika delimiter tidak ditemukan (Gemini tidak patuh format),
       coba parse otomatis berdasarkan heading markdown
    """
    result = {
        "summary":     "",
        "iv_drip":     "",
        "supplements": "",
        "lifestyle":   "",
        "retest":      "",
        "key_priorities": [],
    }

    # Cek apakah response mengandung delimiter
    has_delimiters = any(d in raw for d in _DELIMITERS.values())

    if has_delimiters:
        result = _parse_by_delimiter(raw, result)
    else:
        # Fallback: Gemini tidak pakai delimiter — coba parse by heading
        result = _parse_by_heading(raw, result)

    # Validasi: semua section terisi, fallback ke full text jika ada yang kosong
    result = _ensure_sections_filled(result, raw)

    return result


def _parse_by_delimiter(raw: str, result: dict) -> dict:
    """Split response berdasarkan ===DELIMITER=== markers.
    
    Fix: Gemini sering menambahkan spasi/newline di sekitar delimiter.
    Normalisasi dulu sebelum split agar match selalu berhasil.
    """
    # Normalisasi: pastikan setiap delimiter berdiri di baris sendiri
    normalized = re.sub(r"[ \t]*(===\w+===)[ \t]*", r"\n\1\n", raw)

    delimiter_pattern = r"(===(?:SUMMARY|IV_DRIP|SUPPLEMENTS|LIFESTYLE|RETEST)===)"
    parts = re.split(delimiter_pattern, normalized)

    current_key = None
    for part in parts:
        part = part.strip()
        if not part:
            continue

        matched_key = None
        for key, delimiter in _DELIMITERS.items():
            if part == delimiter:
                matched_key = key
                break

        if matched_key:
            current_key = matched_key
        elif current_key and part:
            # Akumulasi — jangan overwrite jika sudah ada konten
            if result[current_key]:
                result[current_key] += "\n" + part
            else:
                result[current_key] = part

    return result


def _parse_by_heading(raw: str, result: dict) -> dict:
    """
    Fallback parser jika Gemini tidak pakai delimiter.
    Coba deteksi section berdasarkan heading markdown yang umum.
    """
    # Mapping heading keywords ke section keys
    heading_map = {
        "summary":     ["ringkasan", "temuan utama", "summary"],
        "iv_drip":     ["iv drip", "iv therapy", "drip", "infus", "rekomendasi iv"],
        "supplements": ["suplemen", "supplement", "suplementasi"],
        "lifestyle":   ["lifestyle", "gaya hidup", "ice bath", "recovery", "latihan"],
        "retest":      ["retest", "jadwal", "next step", "tindak lanjut"],
    }

    # Split berdasarkan markdown headings (## atau ###)
    sections = re.split(r"\n(?=#{1,3}\s)", raw)

    for section in sections:
        if not section.strip():
            continue

        section_lower = section.lower()
        matched_key = None

        for key, keywords in heading_map.items():
            if any(kw in section_lower for kw in keywords):
                matched_key = key
                break

        if matched_key and not result[matched_key]:
            result[matched_key] = section.strip()

    return result


def _ensure_sections_filled(result: dict, raw_full: str) -> dict:
    """
    Jika masih ada section kosong setelah parsing,
    cek apakah ini karena response terpotong atau format tidak dikenali.
    Sebagai last resort, masukkan full response ke iv_drip agar
    tidak ada konten yang hilang sama sekali.
    """
    required = ["summary", "iv_drip", "supplements", "lifestyle", "retest"]
    empty_sections = [k for k in required if not result.get(k, "").strip()]

    if len(empty_sections) == len(required):
        # Semua kosong — berarti parsing total gagal, dump full response
        result["iv_drip"] = raw_full
        result["summary"] = (
            "Protocol berhasil dibuat. Format response tidak terstruktur — "
            "semua konten ditampilkan di bawah."
        )
    elif empty_sections:
        # Summary boleh kosong — tidak tampilkan pesan error yang membingungkan.
        # Section lain yang kosong: beri pesan singkat saja.
        for key in empty_sections:
            if key == "summary":
                result[key] = ""   # kosong = info box tidak ditampilkan
            else:
                result[key] = "*Konten tidak terbaca. Silakan generate ulang.*"

    return result


def _format_biomarkers_for_prompt(biomarkers: list[dict]) -> str:
    """Format biomarker list menjadi teks ringkas dan prioritas untuk prompt."""
    lines = []

    danger  = [b for b in biomarkers if b.get("status") == "danger"]
    warning = [b for b in biomarkers if b.get("status") == "warning"]
    optimal = [b for b in biomarkers if b.get("status") == "optimal"]

    if danger:
        lines.append("🔴 PERLU PERHATIAN SEGERA:")
        for b in danger:
            ref = f" (Human+ target: {b['reference']})" if b.get("reference") else ""
            lines.append(f"  - {b['name']}: {b['value']} {b['unit']}{ref}")

    if warning:
        lines.append("\n🟡 SUB-OPTIMAL (perlu dioptimalkan):")
        for b in warning:
            ref = f" (Human+ target: {b['reference']})" if b.get("reference") else ""
            lines.append(f"  - {b['name']}: {b['value']} {b['unit']}{ref}")

    if optimal:
        lines.append("\n🟢 SUDAH OPTIMAL (maintain):")
        for b in optimal:
            lines.append(f"  - {b['name']}: {b['value']} {b['unit']}")

    if not lines:
        return "Tidak ada biomarker yang berhasil diekstrak."

    return "\n".join(lines)