"""Aphaea: a cited, cycle-aware wellness education assistant."""

import json
import os
from collections.abc import Iterable
from dataclasses import dataclass
from datetime import date, datetime
from functools import lru_cache
from pathlib import Path

import gradio as gr
import numpy as np
from dotenv import load_dotenv
from huggingface_hub import InferenceClient
from sentence_transformers import SentenceTransformer

APP_DIR = Path(__file__).parent
load_dotenv(APP_DIR / ".env")
MODEL_ID = os.getenv("HF_MODEL", "Qwen/Qwen2.5-7B-Instruct")
HF_TOKEN = os.getenv("HF_TOKEN")
EMBEDDING_MODEL = "sentence-transformers/all-MiniLM-L6-v2"

MEDICAL_NOTICE = (
    "Aphaea provides general education, not diagnosis or medical care. "
    "Seek urgent help for severe pain, fainting, very heavy bleeding, trouble breathing, "
    "or another emergency."
)

SYSTEM_PROMPT = f"""
You are Aphaea, a careful educational assistant for menstrual-cycle wellness,
exercise, contraception, and reproductive health.

Rules:
- Use only the supplied source passages for factual health claims.
- Cite supporting passages inline as [S1], [S2], and so on.
- If the passages do not support an answer, say that clearly and recommend a
  qualified clinician or pharmacist instead of guessing.
- Never diagnose, prescribe, or tell someone to start or stop medication.
- Do not claim that a calendar estimate can confirm ovulation or cycle phase.
- For urgent warning signs, encourage immediate local medical care.
- Keep responses concise, inclusive, respectful, and easy to understand.
- Decline unrelated requests and redirect to Aphaea's health-education scope.

Medical notice: {MEDICAL_NOTICE}
""".strip()


@dataclass(frozen=True)
class SourceChunk:
    text: str
    title: str
    source_id: str
    url: str = ""


def split_legacy_knowledge(text: str) -> list[SourceChunk]:
    """Split the original corpus while retaining stable source identifiers."""
    paragraphs = [p.strip() for p in text.split("\n\n") if len(p.strip()) >= 40]
    return [
        SourceChunk(
            text=paragraph,
            title="Aphaea legacy knowledge base (source review pending)",
            source_id=f"legacy-{index:03d}",
        )
        for index, paragraph in enumerate(paragraphs, start=1)
    ]


class KnowledgeBase:
    def __init__(self, chunks: Iterable[SourceChunk]):
        self.chunks = list(chunks)
        if not self.chunks:
            raise ValueError("The knowledge base contains no usable passages.")
        self.encoder = SentenceTransformer(EMBEDDING_MODEL)
        embeddings = self.encoder.encode(
            [chunk.text for chunk in self.chunks],
            convert_to_numpy=True,
            normalize_embeddings=True,
        )
        self.embeddings = np.asarray(embeddings, dtype=np.float32)

    def search(self, query: str, top_k: int = 4) -> list[tuple[SourceChunk, float]]:
        query_embedding = self.encoder.encode(
            [query], convert_to_numpy=True, normalize_embeddings=True
        )[0]
        scores = self.embeddings @ query_embedding
        count = min(top_k, len(self.chunks))
        indices = np.argsort(scores)[-count:][::-1]
        return [(self.chunks[index], float(scores[index])) for index in indices]


@lru_cache(maxsize=1)
def get_knowledge_base() -> KnowledgeBase:
    records = json.loads(
        (APP_DIR / "knowledge" / "sources.json").read_text(encoding="utf-8")
    )
    chunks = [
        SourceChunk(
            text=record["text"],
            title=record["title"],
            source_id=record["source_id"],
            url=record["url"],
        )
        for record in records
    ]
    return KnowledgeBase(chunks)


def parse_iso_date(value: str | date | datetime) -> date:
    if isinstance(value, datetime):
        parsed = value.date()
    elif isinstance(value, date):
        parsed = value
    elif not value or not str(value).strip():
        raise ValueError("Enter the first day of your most recent period.")
    else:
        try:
            parsed = date.fromisoformat(str(value).strip())
        except ValueError as exc:
            raise ValueError("Choose a valid date from the calendar.") from exc
    if parsed > datetime.now().astimezone().date():
        raise ValueError("The start date cannot be in the future.")
    return parsed


def cycle_day_estimate(start_date: str, typical_length: int) -> tuple[int, str]:
    """Return a transparent calendar estimate, not an ovulation prediction."""
    start = parse_iso_date(start_date)
    length = int(typical_length)
    if not 21 <= length <= 40:
        raise ValueError("Typical cycle length must be between 21 and 40 days.")
    elapsed = (datetime.now().astimezone().date() - start).days
    day = elapsed % length + 1
    return day, (
        f"Calendar estimate: cycle day {day} of a typical {length}-day cycle. "
        "This cannot confirm hormonal phase or ovulation, especially when cycles vary."
    )


def build_profile(start_date: str, typical_length: int, training_level: str):
    try:
        day, estimate = cycle_day_estimate(start_date, typical_length)
    except ValueError as exc:
        return gr.update(value=f"**Check your information:** {exc}"), {}

    profile = {
        "cycle_day": day,
        "cycle_length": int(typical_length),
        "training_level": training_level or "Not specified",
    }
    message = (
        f"### Your private session profile\n\n{estimate}\n\n"
        f"**Training level:** {profile['training_level']}\n\n"
        "Your entries stay in this browser session and are not written to a database."
    )
    return gr.update(value=message), profile


def format_sources(results: list[tuple[SourceChunk, float]]) -> tuple[str, str]:
    context_parts = []
    source_lines = []
    for position, (chunk, score) in enumerate(results, start=1):
        context_parts.append(f"[S{position}] {chunk.text}")
        title = f"[{chunk.title}]({chunk.url})" if chunk.url else chunk.title
        source_lines.append(f"**S{position}.** {title}")
    return "\n\n".join(context_parts), "\n\n".join(source_lines)


def normalize_history(history) -> list[dict[str, str]]:
    normalized = []
    for entry in history or []:
        if isinstance(entry, dict) and entry.get("role") in {"user", "assistant"}:
            content = entry.get("content")
            if isinstance(content, str):
                normalized.append({"role": entry["role"], "content": content})
    return normalized[-8:]


def answer(message: str, history, profile: dict):
    message = (message or "").strip()
    clean_history = normalize_history(history)
    if not message:
        return clean_history, clean_history, "", ""

    try:
        results = get_knowledge_base().search(message)
        context, sources = format_sources(results)
    except Exception as exc:  # noqa: BLE001 - user-facing retrieval boundary
        reply = (
            "I couldn't search the knowledge library right now. "
            f"Please try again after the app is restarted. ({type(exc).__name__})"
        )
        updated = clean_history + [
            {"role": "user", "content": message},
            {"role": "assistant", "content": reply},
        ]
        return updated, updated, "", ""

    if not HF_TOKEN:
        reply = (
            "The interface and knowledge retrieval are working, but model generation "
            "is not configured locally yet. Add `HF_TOKEN` to a local `.env` workflow "
            "or as a Hugging Face Space secret. Never paste the token into source code."
        )
    else:
        profile_note = (
            f"Session profile: {profile}. Treat the cycle day as a calendar estimate only."
            if profile
            else "No optional session profile was provided."
        )
        messages = [
            {"role": "system", "content": SYSTEM_PROMPT},
            *clean_history,
            {
                "role": "user",
                "content": (
                    f"{profile_note}\n\nQuestion: {message}\n\n"
                    f"Retrieved passages:\n{context}"
                ),
            },
        ]
        try:
            client = InferenceClient(model=MODEL_ID, token=HF_TOKEN)
            completion = client.chat_completion(
                messages=messages,
                max_tokens=550,
                temperature=0.2,
            )
            reply = completion.choices[0].message.content or "No response was returned."
        except Exception as exc:  # noqa: BLE001 - remote model boundary
            reply = (
                "The model service is unavailable right now, but your question was not lost. "
                f"Please try again shortly. ({type(exc).__name__})"
            )

    updated = clean_history + [
        {"role": "user", "content": message},
        {"role": "assistant", "content": reply},
    ]
    return updated, updated, "", sources


def suggested_prompt(prompt: str, history, profile):
    return answer(prompt, history, profile)


CSS = """
:root {
  --page: #fbf8f6; --surface: #ffffff; --surface-soft: #f9ecef;
  --ink: #24212b; --muted: #6f6878; --plum: #69405f; --rose: #b75b73;
  --border: #eadfe2; --notice: #fff3f5; --notice-ink: #5d3546;
  --secondary: #f4eff1; --shadow: rgba(69,43,62,.06);
}
.dark {
  --page: #141118; --surface: #211c25; --surface-soft: #30232d;
  --ink: #f7f1f5; --muted: #c8bdc5; --plum: #f0b8cf; --rose: #ef8fab;
  --border: #493a46; --notice: #3a202b; --notice-ink: #ffd9e4;
  --secondary: #332b36; --shadow: rgba(0,0,0,.28);
}
.gradio-container { background: var(--page) !important; color: var(--ink) !important;
  --body-background-fill: var(--page); --block-background-fill: var(--surface);
  --input-background-fill: var(--surface); --body-text-color: var(--ink);
  --block-label-text-color: var(--plum); }
.dark .gradio-container { color-scheme: dark; }
.app-shell { max-width: 1180px; margin: 0 auto; }
.hero { padding: 14px 20px; border: 1px solid var(--border); border-radius: 18px;
  background: linear-gradient(135deg, var(--surface) 0%, var(--surface-soft) 100%); margin-bottom: 10px; }
.eyebrow { color: var(--rose); font-weight: 700; letter-spacing: .12em; text-transform: uppercase; font-size: 12px; }
.hero h1 { font-size: 28px; line-height: 1.05; margin: 3px 0; color: var(--ink); }
.hero p { color: var(--muted); max-width: 850px; font-size: 14px; line-height: 1.3; margin: 0; }
.status { display: inline-block; margin-top: 7px; padding: 4px 9px; border-radius: 999px;
  background: var(--secondary); color: var(--plum); font-size: 12px; font-weight: 700; }
.panel { border: 1px solid var(--border) !important; border-radius: 16px !important;
  background: var(--surface) !important; padding: 12px !important; gap: 8px !important;
  box-shadow: 0 8px 24px var(--shadow); }
.panel h3 { margin: 0 !important; line-height: 1.2 !important; }
.panel button { min-height: 34px !important; }
h1, h2, h3, p { color: var(--ink) !important; }
.panel .block, .panel .wrap, .panel .form { background: var(--surface) !important; color: var(--ink) !important; }
.panel input, .panel textarea { background: var(--surface) !important; color: var(--ink) !important;
  border-color: var(--border) !important; }
.panel input::placeholder, .panel textarea::placeholder { color: var(--muted) !important; opacity: 1; }
#aphaea-chat,
#aphaea-chat .wrap,
#aphaea-chat .message,
#aphaea-chat .message * {
  background: var(--surface) !important;
  color: var(--ink) !important;
}
.dark #aphaea-chat .message.bot, .dark #aphaea-chat .message.bot * { background: #2b2330 !important; }
.panel label, .panel label span, .panel .label-wrap { color: var(--ink) !important; }
button.primary { color: #fff !important; }
.panel button:not(.primary) { background: var(--secondary) !important; color: var(--ink) !important;
  border-color: var(--border) !important; }
.profile-copy { color: var(--muted); font-size: 12px; line-height: 1.35; }
.notice { border-left: 4px solid var(--rose); padding: 7px 9px; background: var(--notice);
  color: var(--notice-ink); border-radius: 8px; font-size: 12px; line-height: 1.3; }
#question-input textarea { min-height: 38px !important; padding-top: 7px !important; padding-bottom: 7px !important; }
.sources { font-size: 12px; color: var(--muted); }
footer { display: none !important; }
@media (max-width: 760px) {
  .hero { padding: 13px 15px; }
  .hero p { font-size: 13px; }
  .panel { padding: 11px !important; }
}
"""

THEME = gr.themes.Base(primary_hue="rose", neutral_hue="slate").set(
    body_background_fill="#fbf8f6",
    body_text_color="#24212b",
    block_background_fill="#ffffff",
    block_border_color="#eadfe2",
    block_label_background_fill="#ffffff",
    block_label_text_color="#69405f",
    input_background_fill="#ffffff",
    input_border_color="#ded6dc",
    input_placeholder_color="#8a838d",
    button_primary_background_fill="#7d4667",
    button_primary_background_fill_hover="#673653",
    button_primary_text_color="#ffffff",
    button_secondary_background_fill="#f4eff1",
    button_secondary_background_fill_hover="#eadfe3",
    button_secondary_text_color="#3f283b",
)


with gr.Blocks(css=CSS, theme=THEME) as demo:
    profile_state = gr.State({})
    history_state = gr.State([])

    with gr.Column(elem_classes="app-shell"):
        gr.HTML(
            """
            <section class="hero">
              <div class="eyebrow">Cycle-aware health education</div>
              <h1>Aphaea</h1>
              <p>Explore menstrual wellness, movement, and contraception with
              answers grounded in a cited source library under active review.</p>
              <span class="status">Educational tool - not medical care</span>
            </section>
            """
        )

        with gr.Row(equal_height=False):
            with gr.Column(scale=1, elem_classes="panel"):
                gr.Markdown("### Optional session profile")
                gr.Markdown(
                    "A calendar estimate can personalize wording, but it cannot confirm "
                    "ovulation or hormonal phase.",
                    elem_classes="profile-copy",
                )
                start_input = gr.DateTime(
                    label="Most recent period start",
                    include_time=False,
                    type="string",
                    info="Choose the first day from the calendar.",
                )
                length_input = gr.Slider(
                    minimum=21,
                    maximum=40,
                    value=28,
                    step=1,
                    label="Typical cycle length",
                )
                training_input = gr.Dropdown(
                    ["Beginner", "Intermediate", "Advanced"],
                    value="Beginner",
                    label="Training experience",
                )
                save_profile = gr.Button("Create session profile", variant="primary")
                profile_output = gr.Markdown(
                    "No profile created. You can still use the assistant."
                )
                gr.HTML(f'<div class="notice">{MEDICAL_NOTICE}</div>')

            with gr.Column(scale=2, elem_classes="panel"):
                gr.Markdown("### Ask Aphaea")
                chatbot = gr.Chatbot(
                    elem_id="aphaea-chat",
                    type="messages",
                    height=210,
                    placeholder="Ask a question or choose a starting prompt below.",
                    show_copy_button=True,
                )
                user_input = gr.Textbox(
                    elem_id="question-input",
                    label="Your question",
                    placeholder="For example: What are common reasons a cycle changes?",
                    lines=1,
                )
                with gr.Row():
                    send_button = gr.Button("Send", variant="primary")
                    clear_button = gr.Button("Clear conversation")
                with gr.Row():
                    workout_button = gr.Button("Adapt a workout")
                    contraception_button = gr.Button("Compare contraception")
                    symptoms_button = gr.Button("Understand cycle changes")
                with gr.Accordion("Sources used for the latest answer", open=False):
                    source_output = gr.Markdown(
                        "Sources will appear after you ask a question.",
                        elem_classes="sources",
                    )

    save_profile.click(
        build_profile,
        [start_input, length_input, training_input],
        [profile_output, profile_state],
    )

    for trigger in (send_button.click, user_input.submit):
        trigger(
            answer,
            [user_input, history_state, profile_state],
            [chatbot, history_state, user_input, source_output],
        )

    workout_button.click(
        lambda history, profile: suggested_prompt(
            "How can I adjust exercise when my energy or symptoms change?",
            history,
            profile,
        ),
        [history_state, profile_state],
        [chatbot, history_state, user_input, source_output],
    )
    contraception_button.click(
        lambda history, profile: suggested_prompt(
            "What factors should someone discuss with a clinician when comparing contraception?",
            history,
            profile,
        ),
        [history_state, profile_state],
        [chatbot, history_state, user_input, source_output],
    )
    symptoms_button.click(
        lambda history, profile: suggested_prompt(
            "What are common reasons menstrual-cycle timing can change?",
            history,
            profile,
        ),
        [history_state, profile_state],
        [chatbot, history_state, user_input, source_output],
    )
    clear_button.click(
        lambda: ([], [], "Sources will appear after you ask a question."),
        outputs=[chatbot, history_state, source_output],
    )


if __name__ == "__main__":
    demo.queue().launch()