"""
Desensitization pipeline for enterprise data protection.

Pure function library -- no class state. LLM client is passed in.

Core flow: input(List[Dict]) -> output(str without commentary)
"""
from __future__ import annotations

import hashlib
import json
import logging
from typing import Any, Dict, List, Optional, Set

from actor.connectors.llm import LLMClient, TOKEN_LIMIT_CASCADE

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Prompts
# ---------------------------------------------------------------------------

FIELD_CLASSIFICATION_PROMPT = """\
You are a data classification engine. Given field names, their descriptions, and sample values from a database table, classify each field as either "sensitive" or "safe".

Sensitive fields include: personal names, email addresses, phone numbers, social security numbers, financial data (salaries, account numbers, credit cards), medical information, passwords, tokens, addresses, dates of birth, and any other personally identifiable information (PII).

Safe fields include: department names, status codes, categories, timestamps of business events, non-identifying IDs, and other non-PII business data.

Table context:
{table_context}

Fields to classify (with sample values):
{field_samples}

Entity ID fields (always safe, used for joining): {entity_ids}

Respond with ONLY a JSON object mapping each field name to "sensitive" or "safe". No commentary.
Example: {{"name": "sensitive", "department": "safe", "email": "sensitive"}}
"""

DESENSITIZE_CHUNK_PROMPT = """\
You are a data desensitization engine. Transform the following pre-masked database rows into a safe textual representation. Sensitive fields have already been replaced with hash tokens (e.g., NAM_a1b2c3d4). Preserve the hash tokens exactly as they appear. Convert the data into a clean, readable text format suitable for analytics -- one record per line, fields separated by pipes.

Return ONLY the transformed text. No commentary, no explanations, no markdown formatting.

Data:
{data}
"""


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _hash_value(value: Any, field: str) -> str:
    """Deterministic SHA256-based hash with type prefix.

    Produces tokens like NAM_a1b2c3d4 for consistent masking across chunks.
    """
    prefix_map = {
        "name": "NAM", "email": "EML", "phone": "PHN", "address": "ADR",
        "salary": "FIN", "ssn": "SSN", "dob": "DOB", "account": "ACC",
    }
    # Determine prefix from field name heuristics
    field_lower = field.lower()
    prefix = "MSK"
    for key, pfx in prefix_map.items():
        if key in field_lower:
            prefix = pfx
            break
    digest = hashlib.sha256(f"{field}:{value}".encode()).hexdigest()[:8]
    return f"{prefix}_{digest}"


def _estimate_tokens(text: str, model: str = "") -> int:
    """Estimate token count. Uses tiktoken if available, else len/4 fallback."""
    try:
        import tiktoken
        enc = tiktoken.encoding_for_model(model) if model else tiktoken.get_encoding("cl100k_base")
        return len(enc.encode(text))
    except Exception:
        return len(text) // 4


def _chunk_data(data: List[Dict[str, Any]], max_tokens: int, model: str = "") -> List[List[Dict[str, Any]]]:
    """Split rows into token-bounded chunks."""
    if not data:
        return []
    chunks: List[List[Dict[str, Any]]] = []
    current_chunk: List[Dict[str, Any]] = []
    current_tokens = 0

    for row in data:
        row_text = json.dumps(row)
        row_tokens = _estimate_tokens(row_text, model)
        if current_chunk and (current_tokens + row_tokens) > max_tokens:
            chunks.append(current_chunk)
            current_chunk = []
            current_tokens = 0
        current_chunk.append(row)
        current_tokens += row_tokens

    if current_chunk:
        chunks.append(current_chunk)
    return chunks


def _get_sample_values(data: List[Dict[str, Any]], max_samples: int = 3) -> Dict[str, List[str]]:
    """Extract sample values per field for classification."""
    samples: Dict[str, List[str]] = {}
    for row in data[:max_samples * 2]:  # scan a few extra rows for variety
        for field, value in row.items():
            if field not in samples:
                samples[field] = []
            if value is not None and len(samples[field]) < max_samples:
                samples[field].append(str(value))
    return samples


def _classify_fields(
    data: List[Dict[str, Any]],
    table_context: Dict[str, str],
    entity_ids: List[str],
    llm_client: LLMClient,
    model: str,
) -> Dict[str, str]:
    """Classify fields as sensitive or safe using LLM.

    Uses table/field descriptions from SaaS + sample values from data.
    """
    if not data:
        return {}

    samples = _get_sample_values(data)
    field_samples_parts = []
    for field, vals in samples.items():
        desc = table_context.get(field, "no description")
        field_samples_parts.append(f"  {field} (desc: {desc}): {vals}")
    field_samples_str = "\n".join(field_samples_parts)

    table_desc = table_context.get("table", "no table description")
    context_str = f"Table: {table_desc}"

    prompt = FIELD_CLASSIFICATION_PROMPT.format(
        table_context=context_str,
        field_samples=field_samples_str,
        entity_ids=entity_ids,
    )

    try:
        result = llm_client.completion(
            messages=[{"role": "user", "content": prompt}],
            temperature=0.0,
            max_tokens=1024,
        )
        content = result.get("content", "")
        # Strip markdown code fences if present
        content = content.strip()
        if content.startswith("```"):
            lines = content.split("\n")
            lines = [l for l in lines if not l.strip().startswith("```")]
            content = "\n".join(lines)
        classification = json.loads(content)
        # Ensure entity_ids are always "safe"
        for eid in entity_ids:
            classification[eid] = "safe"
        return classification
    except Exception as exc:
        logger.warning("Field classification failed: %s; defaulting all to sensitive", exc)
        fields = set()
        for row in data:
            fields.update(row.keys())
        classification = {f: "sensitive" for f in fields}
        for eid in entity_ids:
            classification[eid] = "safe"
        return classification


def _build_masking_map(
    data: List[Dict[str, Any]],
    classification: Dict[str, str],
    entity_ids: List[str],
) -> Dict[str, Dict[str, str]]:
    """Build deterministic mask map for sensitive fields.

    Returns: {field_name: {original_value: masked_value}}
    """
    entity_id_set: Set[str] = set(entity_ids)
    mask_map: Dict[str, Dict[str, str]] = {}
    for row in data:
        for field, value in row.items():
            if field in entity_id_set:
                continue
            if classification.get(field) == "sensitive" and value is not None:
                if field not in mask_map:
                    mask_map[field] = {}
                str_val = str(value)
                if str_val not in mask_map[field]:
                    mask_map[field][str_val] = _hash_value(value, field)
    return mask_map


def _apply_masking(
    data: List[Dict[str, Any]],
    mask_map: Dict[str, Dict[str, str]],
) -> List[Dict[str, Any]]:
    """Apply masks to rows, returning new list of dicts."""
    masked: List[Dict[str, Any]] = []
    for row in data:
        new_row: Dict[str, Any] = {}
        for field, value in row.items():
            if field in mask_map and value is not None:
                str_val = str(value)
                new_row[field] = mask_map[field].get(str_val, str_val)
            else:
                new_row[field] = value
        masked.append(new_row)
    return masked


def _process_chunk(
    chunk: List[Dict[str, Any]],
    mask_map: Dict[str, Dict[str, str]],
    entity_ids: List[str],
    llm_client: LLMClient,
    model: str,
) -> str:
    """Process a single chunk through masking + LLM, returns string."""
    masked = _apply_masking(chunk, mask_map)
    data_str = json.dumps(masked, default=str)
    prompt = DESENSITIZE_CHUNK_PROMPT.format(data=data_str)

    result = llm_client.completion(
        messages=[{"role": "user", "content": prompt}],
        temperature=0.0,
    )
    return result.get("content", "").strip()


def desensitize(
    data: List[Dict[str, Any]],
    table_context: Dict[str, str],
    entity_ids: List[str],
    llm_client: LLMClient,
    model: str,
    chunk_max_tokens: int = 4096,
) -> str:
    """Full desensitization pipeline.

    1. Classify fields as sensitive/safe using LLM + table context
    2. Build deterministic masking map for sensitive values
    3. Chunk data by token limit
    4. Process each chunk through LLM with pre-masked data
    5. Return concatenated desensitized text

    On token-exceeded error from LLM, scales down via TOKEN_LIMIT_CASCADE and retries.
    """
    if not data:
        return ""

    # Step 1: Classify
    classification = _classify_fields(data, table_context, entity_ids, llm_client, model)

    # Step 2: Build mask map
    mask_map = _build_masking_map(data, classification, entity_ids)

    # Step 3: Chunk
    current_limit = chunk_max_tokens
    chunks = _chunk_data(data, current_limit, model)

    # Step 4: Process chunks with token cascade retry
    parts: List[str] = []
    for chunk in chunks:
        retry_limit = current_limit
        while True:
            try:
                part = _process_chunk(chunk, mask_map, entity_ids, llm_client, model)
                parts.append(part)
                break
            except Exception as exc:
                exc_str = str(exc).lower()
                if "token" in exc_str and ("limit" in exc_str or "exceed" in exc_str or "maximum" in exc_str):
                    new_limit = LLMClient.scale_down_token_limit(retry_limit)
                    if new_limit >= retry_limit:
                        raise
                    logger.warning(
                        "Token limit exceeded; scaling down from %d to %d",
                        retry_limit, new_limit,
                    )
                    retry_limit = new_limit
                    # Re-chunk this specific chunk with smaller limit
                    sub_chunks = _chunk_data(chunk, retry_limit, model)
                    for sc in sub_chunks:
                        sub_part = _process_chunk(sc, mask_map, entity_ids, llm_client, model)
                        parts.append(sub_part)
                    break
                else:
                    raise

    return "\n".join(parts)
