"""
Context processing utilities for Self Vision POC.

Handles:
- Text normalization (strip nulls, empty blocks)
- Word-based truncation (keep last N words)
- Combined LLM input construction with section headers
"""

import re


def normalize_context(text: str) -> str:
    """
    Normalize raw context text.

    Removes:
        - Literal 'null', 'undefined', 'None' tokens
        - Empty lines / whitespace-only blocks

    Preserves:
        - Timestamps
        - Speaker labels
        - Original ordering
    """
    if not text:
        return ""

    # Remove literal null/undefined/None tokens (standalone words)
    cleaned = re.sub(r'\b(null|undefined|None)\b', '', text)

    # Remove lines that are entirely whitespace after cleaning
    lines = cleaned.split('\n')
    non_empty = [line for line in lines if line.strip()]

    result = '\n'.join(non_empty)

    # Collapse runs of 3+ whitespace into double newline
    result = re.sub(r'\n{3,}', '\n\n', result)

    return result.strip()


def truncate_last_words(text: str, max_words: int = 2000) -> str:
    """
    Keep only the last `max_words` words from `text`.

    Word-based truncation only — never character-wise.
    """
    if not text:
        return ""

    words = text.split()
    if len(words) <= max_words:
        return text

    # Keep the last max_words words
    truncated_words = words[-max_words:]
    return ' '.join(truncated_words)


def count_words(text: str) -> int:
    """Count words in text."""
    if not text:
        return 0
    return len(text.split())


def normalize_roadmap(text: str) -> str:
    """
    Format roadmap text. If it is JSON, pretty-print it.
    Otherwise, fallback to standard context normalization.
    """
    if not text:
        return ""
    import json
    try:
        data = json.loads(text)
        return json.dumps(data, indent=2)
    except Exception:
        return normalize_context(text)


def build_llm_input(
    onboarding_text: str,
    transcript_text: str,
    roadmap_text: str = "",
    lesson_context: str = "",
    max_total_words: int = 6000,
    per_section_limit: int = 2000,
) -> dict:
    """
    Build the combined LLM input from onboarding context, roadmap context, and current transcript.

    Ordering: older context first, newest context last.

    If combined word count > max_total_words, each section is independently
    truncated to keep the last `per_section_limit` words.

    Returns:
        dict with keys:
            - 'combined_text': the final prompt text
            - 'total_word_count': word count of combined text
            - 'truncation_applied': bool
    """
    onboarding_clean = normalize_context(onboarding_text)
    transcript_clean = normalize_context(transcript_text)
    roadmap_clean = normalize_roadmap(roadmap_text)

    total_words = (
        count_words(onboarding_clean)
        + count_words(transcript_clean)
        + count_words(roadmap_clean)
    )
    truncation_applied = False

    if total_words > max_total_words:
        onboarding_clean = truncate_last_words(onboarding_clean, per_section_limit)
        transcript_clean = truncate_last_words(transcript_clean, per_section_limit)
        roadmap_clean = truncate_last_words(roadmap_clean, per_section_limit)
        truncation_applied = True

    # Build combined text with section headers
    sections = []

    sections.append("[ONBOARDING_CONTEXT]")
    sections.append("")
    sections.append(onboarding_clean if onboarding_clean else "(no prior onboarding context)")
    sections.append("")

    sections.append("[ROADMAP_CONTEXT]")
    sections.append("")
    sections.append(roadmap_clean if roadmap_clean else "(no prior career roadmap context)")
    sections.append("")

    sections.append("[SELF_VISION_TRANSCRIPT]")
    sections.append("")
    sections.append(transcript_clean)

    if lesson_context and lesson_context.strip():
        sections.append("")
        sections.append("[LESSON_CONTEXT]")
        sections.append("")
        sections.append(normalize_context(lesson_context))

    combined = '\n'.join(sections)
    final_word_count = count_words(combined)

    return {
        "combined_text": combined,
        "total_word_count": final_word_count,
        "truncation_applied": truncation_applied,
    }