"""
Adaptive context engine.
Tracks user session history, infers expertise level, and generates
context-aware system prompt extensions for the custom GPT.
"""
from __future__ import annotations
import time
from collections import defaultdict
from dataclasses import dataclass, field
from typing import Any

# ── In-memory session store (per HF Space instance) ───────────────────────────

@dataclass
class Session:
    session_id: str
    created_at: float = field(default_factory=time.time)
    proteins_visited: list[str] = field(default_factory=list)
    question_complexity: list[int] = field(default_factory=list)
    domains_of_interest: list[str] = field(default_factory=list)
    interaction_count: int = 0
    expertise_score: float = 0.5         # 0=novice, 1=expert
    preferred_depth: str = "balanced"    # brief | balanced | deep

_sessions: dict[str, Session] = {}


def get_session(session_id: str) -> Session:
    if session_id not in _sessions:
        _sessions[session_id] = Session(session_id=session_id)
    return _sessions[session_id]


# ── Expertise inference ────────────────────────────────────────────────────────

EXPERT_TERMS = {
    "plddt", "alphafold", "uniprot", "pfam", "ptm", "idr", "cryo-em",
    "x-ray crystallography", "ramachandran", "secondary structure",
    "disordered region", "binding affinity", "allosteric", "catalytic triad",
    "homology", "ortholog", "paralog", "domain", "residue", "motif",
    "b-factor", "resolution", "rmsd", "torsion angle",
}

NOVICE_TERMS = {
    "what is", "explain", "how does", "simple", "basic", "beginner",
    "tell me about", "meaning of", "what does",
}

def _infer_complexity(query: str) -> int:
    """Return 0 (novice) to 2 (expert) based on query vocabulary."""
    q = query.lower()
    expert_hits = sum(1 for t in EXPERT_TERMS if t in q)
    novice_hits = sum(1 for t in NOVICE_TERMS if t in q)
    if expert_hits >= 2:
        return 2
    if expert_hits == 1:
        return 1
    if novice_hits > 0:
        return 0
    return 1


# ── Adaptive system prompt builder ────────────────────────────────────────────

DEPTH_INSTRUCTIONS = {
    "brief": (
        "Keep responses concise. Use plain language. "
        "Avoid jargon. Focus on practical takeaways."
    ),
    "balanced": (
        "Balance scientific accuracy with clarity. "
        "Define technical terms on first use. "
        "Provide one concrete example per concept."
    ),
    "deep": (
        "Assume expert-level structural biology knowledge. "
        "Include mechanistic detail, cite pLDDT thresholds, "
        "reference domain databases (Pfam, InterPro), and discuss "
        "functional implications at the residue level."
    ),
}

def build_adaptive_prompt(session: Session, current_query: str = "") -> str:
    """Return a context block that the GPT prepends to its system prompt."""
    if current_query:
        c = _infer_complexity(current_query)
        session.question_complexity.append(c)
        session.interaction_count += 1
        # Exponential moving average for expertise score
        alpha = 0.3
        session.expertise_score = (
            alpha * (c / 2.0) + (1 - alpha) * session.expertise_score
        )

    # Decide depth
    if session.expertise_score > 0.65:
        session.preferred_depth = "deep"
    elif session.expertise_score < 0.35:
        session.preferred_depth = "brief"
    else:
        session.preferred_depth = "balanced"

    depth_instr = DEPTH_INSTRUCTIONS[session.preferred_depth]

    # Build context block
    history_line = ""
    if session.proteins_visited:
        recent = session.proteins_visited[-5:]
        history_line = f"Proteins discussed this session: {', '.join(recent)}."

    domain_line = ""
    if session.domains_of_interest:
        domains = list(set(session.domains_of_interest[-6:]))
        domain_line = f"User's areas of interest: {', '.join(domains)}."

    prompt = f"""[ADAPTIVE CONTEXT — session {session.session_id[:8]}]
Detected expertise level: {session.preferred_depth.upper()} ({session.expertise_score:.2f}/1.0).
Instruction: {depth_instr}
{history_line}
{domain_line}
Interaction count this session: {session.interaction_count}.
Tailor every response to the above profile without explicitly mentioning it."""

    return prompt.strip()


def update_session_protein(session: Session, uniprot_id: str, domains: list[str]) -> None:
    if uniprot_id not in session.proteins_visited:
        session.proteins_visited.append(uniprot_id)
    session.domains_of_interest.extend(domains)


# ── Adaptive follow-up question generator ────────────────────────────────────

_FOLLOWUPS: dict[str, list[str]] = {
    "novice": [
        "What does this protein actually do in the human body?",
        "Is this protein associated with any disease?",
        "What happens if this protein is mutated?",
        "Can you explain what the confidence colours mean?",
    ],
    "balanced": [
        "Which residues form the active site or binding interface?",
        "Are there known pathogenic mutations in this protein?",
        "What structural domains does this protein contain?",
        "How does pLDDT score vary across regions — and what does that imply?",
    ],
    "expert": [
        "What is the allosteric communication network within this structure?",
        "Which disordered regions are functionally important?",
        "How does the predicted structure compare to experimental PDB entries?",
        "What post-translational modifications are predicted at high-pLDDT residues?",
    ],
}

def generate_followups(session: Session, n: int = 3) -> list[str]:
    level = session.preferred_depth
    key = {"brief": "novice", "balanced": "balanced", "deep": "expert"}[level]
    pool = _FOLLOWUPS[key]
    visited = len(session.proteins_visited)
    start = (visited * 2) % len(pool)
    return [pool[(start + i) % len(pool)] for i in range(n)]