

# DomainProfile loads customer domain data for data-dependent attacks.

# Supported entry types (in the JSONL bundle):
#  train_sample     — text seen during fine-tuning (for membership inference)
#   holdout_sample   — text NOT seen during fine-tuning (control set)
#   domain_fact      — Q&A pair for hallucination probes
#  sensitive_string — PII / PHI strings to probe for extraction

# File format: newline-delimited JSON (.jsonl), one entry per line.

# bundle up mannn 
#   {"type": "train_sample",     "text": "Patient ID 4821 was admitted on..."}
#   {"type": "holdout_sample",   "text": "The quarterly report shows..."}
#   {"type": "domain_fact",      "question": "What is the max dosage?", "answer": "500mg"}
#   {"type": "sensitive_string", "text": "SSN: 078-05-1120", "category": "ssn"}

# Usage:
#     profile = DomainProfile.from_jsonl("domain_data.jsonl")
#     train   = profile.train_samples()     # List[str]
#    holdout = profile.holdout_samples()   # List[str]
#     facts   = profile.domain_facts()      # List[DomainFact]
#     pii     = profile.sensitive_strings() # List[SensitiveString]
#     summary = profile.summary()           # dict — safe to log
#test asap!!!!!


from __future__ import annotations

import json
import logging
from dataclasses import dataclass, field
from pathlib import Path
from typing import List, Optional

logger = logging.getLogger(__name__)

VALID_TYPES = {"train_sample", "holdout_sample", "domain_fact", "sensitive_string"}


# ---------------------------------------------------------------------------
# Entry dataclasses
# ---------------------------------------------------------------------------

@dataclass
class TrainSample:
    text: str


@dataclass
class HoldoutSample:
    text: str


@dataclass
class DomainFact:
    question: str
    answer: str
    context: Optional[str] = None  # optional surrounding context


@dataclass
class SensitiveString:
    text: str
    category: str = "unknown"  # e.g. "ssn", "dob", "patient_id", "phone"


# ---------------------------------------------------------------------------
# DomainProfile
# ---------------------------------------------------------------------------

@dataclass
class DomainProfile:
    _train:   List[TrainSample]    = field(default_factory=list)
    _holdout: List[HoldoutSample]  = field(default_factory=list)
    _facts:   List[DomainFact]     = field(default_factory=list)
    _pii:     List[SensitiveString] = field(default_factory=list)
    source_path: Optional[str] = None

    # ------------------------------------------------------------------
    # Loaders
    # ------------------------------------------------------------------

    @classmethod
    def from_jsonl(cls, path: str | Path) -> "DomainProfile":
        path = Path(path)
        if not path.exists():
            raise FileNotFoundError(f"Domain data not found: {path}")

        profile = cls(source_path=str(path))
        skipped = 0

        with open(path, encoding="utf-8") as f:
            for i, line in enumerate(f, 1):
                line = line.strip()
                if not line or line.startswith("#"):
                    continue
                try:
                    entry = json.loads(line)
                    profile._ingest(entry, line_number=i)
                except json.JSONDecodeError as e:
                    logger.warning(f"Line {i}: invalid JSON — {e}")
                    skipped += 1

        if skipped:
            logger.warning(f"Skipped {skipped} malformed lines in {path}")

        logger.info(
            f"DomainProfile loaded from {path}: "
            f"{len(profile._train)} train, {len(profile._holdout)} holdout, "
            f"{len(profile._facts)} facts, {len(profile._pii)} sensitive strings"
        )
        return profile

    @classmethod
    def empty(cls) -> "DomainProfile":
        """Returns an empty profile — attacks that need data will skip gracefully."""
        return cls()

    # ------------------------------------------------------------------
    # Internal
    # ------------------------------------------------------------------

    def _ingest(self, entry: dict, line_number: int):
        t = entry.get("type")
        if t not in VALID_TYPES:
            logger.warning(f"Line {line_number}: unknown type '{t}' — skipping")
            return

        if t == "train_sample":
            text = entry.get("text", "").strip()
            if text:
                self._train.append(TrainSample(text=text))

        elif t == "holdout_sample":
            text = entry.get("text", "").strip()
            if text:
                self._holdout.append(HoldoutSample(text=text))

        elif t == "domain_fact":
            q = entry.get("question", "").strip()
            a = entry.get("answer", "").strip()
            if q and a:
                self._facts.append(DomainFact(
                    question=q,
                    answer=a,
                    context=entry.get("context"),
                ))
            else:
                logger.warning(f"Line {line_number}: domain_fact missing question or answer")

        elif t == "sensitive_string":
            text = entry.get("text", "").strip()
            if text:
                self._pii.append(SensitiveString(
                    text=text,
                    category=entry.get("category", "unknown"),
                ))

    # ------------------------------------------------------------------
    # Public accessors
    # ------------------------------------------------------------------

    def train_samples(self) -> List[str]:
        return [s.text for s in self._train]

    def holdout_samples(self) -> List[str]:
        return [s.text for s in self._holdout]

    def domain_facts(self) -> List[DomainFact]:
        return list(self._facts)

    def sensitive_strings(self) -> List[SensitiveString]:
        return list(self._pii)

    def has_membership_inference_data(self) -> bool:
        return len(self._train) > 0 and len(self._holdout) > 0

    def has_hallucination_data(self) -> bool:
        return len(self._facts) > 0

    def has_extraction_data(self) -> bool:
        return len(self._pii) > 0

    def is_empty(self) -> bool:
        return not (self._train or self._holdout or self._facts or self._pii)

    def summary(self) -> dict:
        """Safe-to-log summary — no actual sensitive content."""
        return {
            "source": self.source_path,
            "train_samples": len(self._train),
            "holdout_samples": len(self._holdout),
            "domain_facts": len(self._facts),
            "sensitive_strings": len(self._pii),
            "capabilities": {
                "membership_inference": self.has_membership_inference_data(),
                "hallucination_probes": self.has_hallucination_data(),
                "extraction_probes": self.has_extraction_data(),
            },
        }