import re
import json

from promptlock.exceptions import ContractViolation


class OutputContract:
    """Defines and validates the expected shape of an LLM output."""

    def __init__(
        self,
        required_fields: list[str] = None,
        max_length: int = None,
        min_length: int = None,
        regex_patterns: list[str] = None,
        allowed_values: list[str] = None,
    ):
        self.required_fields = required_fields or []
        self.max_length = max_length
        self.min_length = min_length
        self.regex_patterns = regex_patterns or []
        self.allowed_values = allowed_values or []

    def validate(self, output: str | dict) -> bool:
        """Validate an LLM output against this contract.

        Args:
            output: the raw LLM output, either a string or a dict (parsed JSON)

        Returns:
            True if all checks pass

        Raises:
            ContractViolation: if any rule is violated
        """
        parsed = self._parse(output)
        text = output if isinstance(output, str) else json.dumps(output)

        self._check_required_fields(parsed)
        self._check_length(text)
        self._check_regex(text)
        self._check_allowed_values(output)

        return True

    def _parse(self, output: str | dict) -> dict | None:
        """Try to parse output as JSON if it's a string."""
        if isinstance(output, dict):
            return output
        try:
            return json.loads(output)
        except (json.JSONDecodeError, TypeError):
            return None

    def _check_required_fields(self, parsed: dict | None):
        if not self.required_fields:
            return
        if parsed is None:
            raise ContractViolation(
                message="Output must be a JSON object to validate required fields.",
                field="required_fields",
            )
        missing = [f for f in self.required_fields if f not in parsed]
        if missing:
            raise ContractViolation(
                message=f"Missing required fields: {missing}",
                field="required_fields",
            )

    def _check_length(self, text: str):
        length = len(text)
        if self.min_length is not None and length < self.min_length:
            raise ContractViolation(
                message=f"Output length {length} is below minimum {self.min_length}.",
                field="min_length",
            )
        if self.max_length is not None and length > self.max_length:
            raise ContractViolation(
                message=f"Output length {length} exceeds maximum {self.max_length}.",
                field="max_length",
            )

    def _check_regex(self, text: str):
        for pattern in self.regex_patterns:
            if not re.search(pattern, text):
                raise ContractViolation(
                    message=f"Output does not match required pattern: '{pattern}'",
                    field="regex_patterns",
                )

    def _check_allowed_values(self, output: str | dict):
        if not self.allowed_values:
            return
        if not isinstance(output, str):
            raise ContractViolation(
                message="allowed_values can only be checked against string outputs.",
                field="allowed_values",
            )
        if output not in self.allowed_values:
            raise ContractViolation(
                message=f"Output '{output}' is not in allowed values: {self.allowed_values}",
                field="allowed_values",
            )

    def __repr__(self) -> str:
        return (
            f"OutputContract("
            f"required_fields={self.required_fields}, "
            f"max_length={self.max_length}, "
            f"min_length={self.min_length}, "
            f"regex_patterns={self.regex_patterns}, "
            f"allowed_values={self.allowed_values})"
        )