"""
Output validation: schema enforcement, hallucination detection, prompt leak detection.
"""

import json
import re
from dataclasses import dataclass, field
from typing import Any, Optional, Union


@dataclass
class ValidationResult:
    valid: bool
    output: str
    warnings: list = field(default_factory=list)
    errors: list = field(default_factory=list)


# System prompt delimiter patterns that shouldn't appear in LLM output
_LEAKED_PROMPT_PATTERNS = [
    re.compile(r"<\s*/?system\s*>", re.IGNORECASE),
    re.compile(r"\[INST\]|\[/INST\]", re.IGNORECASE),
    re.compile(r"<<SYS>>|<</SYS>>", re.IGNORECASE),
    re.compile(r"<\|im_start\|>|<\|im_end\|>"),
    re.compile(r"###\s*(Human|Assistant|System)\s*:", re.IGNORECASE),
    re.compile(r"<\|endoftext\|>"),
    re.compile(r"SYSTEM\s*PROMPT\s*:", re.IGNORECASE),
    re.compile(r"You are (a |an )?(helpful\s+)?assistant.*?\.", re.IGNORECASE | re.DOTALL),
]

# URL pattern for hallucination detection
_URL_PATTERN = re.compile(r"https?://([^\s/\"'<>\])+)([^\s\"'<>\]]*)")

# Known trusted TLD suffixes (non-exhaustive, covers common real domains)
_SUSPICIOUS_URL_SIGNALS = [
    re.compile(r"https?://[a-z0-9-]+\.(example|test|localhost|fake|invalid|placeholder)\b", re.IGNORECASE),
    re.compile(r"https?://[a-z]{20,}\.[a-z]{2,6}\b"),  # suspiciously long random-looking domain
    re.compile(r"https?://\d{4,}\.[a-z]{2,6}\b"),  # numeric-only domain
]


def _check_leaked_prompts(text: str) -> list:
    warnings = []
    for pattern in _LEAKED_PROMPT_PATTERNS:
        if pattern.search(text):
            warnings.append(f"Possible system prompt leak detected (pattern: {pattern.pattern[:40]}...)")
    return warnings


def _check_hallucinated_urls(text: str) -> list:
    warnings = []
    urls = _URL_PATTERN.findall(text)
    for domain, _ in urls:
        for sig_pattern in _SUSPICIOUS_URL_SIGNALS:
            full_url = f"https://{domain}"
            if sig_pattern.match(full_url):
                warnings.append(f"Potentially hallucinated URL detected: https://{domain}")
                break
    return warnings


def _validate_json_schema(text: str, schema: Optional[dict]) -> tuple:
    """
    Try to parse JSON and optionally validate against a dict schema.
    Returns (parsed_obj_or_None, errors_list).
    """
    stripped = text.strip()
    # Try to extract JSON block if wrapped in markdown fences
    fence_match = re.search(r"```(?:json)?\s*([\s\S]+?)```", stripped, re.IGNORECASE)
    if fence_match:
        stripped = fence_match.group(1).strip()

    try:
        obj = json.loads(stripped)
    except json.JSONDecodeError as e:
        return None, [f"JSON parse error: {e}"]

    errors = []
    if schema:
        # Basic key presence validation
        for key, expected_type in schema.items():
            if key not in obj:
                errors.append(f"Missing required key: '{key}'")
            elif expected_type is not None and not isinstance(obj[key], expected_type):
                errors.append(
                    f"Key '{key}' expected {expected_type.__name__}, "
                    f"got {type(obj[key]).__name__}"
                )

    return obj, errors


def enforce(
    text: str,
    schema: Optional[Union[str, dict]] = None,
    check_leaks: bool = True,
    check_urls: bool = True,
) -> ValidationResult:
    """
    Validate LLM output against optional schema, and check for safety signals.

    Args:
        text: Raw LLM output string.
        schema: "json" to require valid JSON, or a dict mapping key→expected_type
                for key-presence + type checking. None for no schema enforcement.
        check_leaks: If True, scan for leaked system prompt delimiters.
        check_urls: If True, flag suspicious / likely-hallucinated URLs.

    Returns:
        ValidationResult with .valid, .output, .warnings, .errors
    """
    warnings = []
    errors = []
    output = text

    # Schema validation
    if schema == "json":
        obj, json_errors = _validate_json_schema(text, None)
        errors.extend(json_errors)
        if obj is not None:
            output = json.dumps(obj)
    elif isinstance(schema, dict):
        obj, json_errors = _validate_json_schema(text, schema)
        errors.extend(json_errors)
        if obj is not None:
            output = json.dumps(obj)

    # Prompt leak detection
    if check_leaks:
        warnings.extend(_check_leaked_prompts(text))

    # URL hallucination detection
    if check_urls:
        warnings.extend(_check_hallucinated_urls(text))

    return ValidationResult(
        valid=len(errors) == 0,
        output=output,
        warnings=warnings,
        errors=errors,
    )
