"""Prompt tester — automated testing for LLM prompts."""

import json
import re
import time
from dataclasses import dataclass, field
from enum import Enum
from pathlib import Path
from typing import Any, Callable, Dict, List, Optional

from promptkit.prompt import Prompt, PromptVersion


class TestStatus(str, Enum):
    """Test result status."""
    PASSED = "passed"
    FAILED = "failed"
    SKIPPED = "skipped"
    ERROR = "error"


@dataclass
class TestCase:
    """A single test case for a prompt."""
    name: str
    description: str = ""
    variables: Dict[str, Any] = field(default_factory=dict)
    assertions: List[Dict[str, Any]] = field(default_factory=list)
    tags: List[str] = field(default_factory=list)

    def to_dict(self) -> Dict[str, Any]:
        return {
            "name": self.name,
            "description": self.description,
            "variables": self.variables,
            "assertions": self.assertions,
            "tags": self.tags,
        }

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "TestCase":
        return cls(
            name=data["name"],
            description=data.get("description", ""),
            variables=data.get("variables", {}),
            assertions=data.get("assertions", []),
            tags=data.get("tags", []),
        )


@dataclass
class TestResult:
    """Result of running a single test case."""
    test_name: str
    status: TestStatus
    message: str = ""
    duration_ms: float = 0.0
    details: Dict[str, Any] = field(default_factory=dict)

    def to_dict(self) -> Dict[str, Any]:
        return {
            "test_name": self.test_name,
            "status": self.status.value,
            "message": self.message,
            "duration_ms": self.duration_ms,
            "details": self.details,
        }


@dataclass
class TestReport:
    """Complete test report."""
    prompt_name: str
    version: int
    results: List[TestResult] = field(default_factory=list)
    total_duration_ms: float = 0.0

    @property
    def passed(self) -> int:
        return sum(1 for r in self.results if r.status == TestStatus.PASSED)

    @property
    def failed(self) -> int:
        return sum(1 for r in self.results if r.status == TestStatus.FAILED)

    @property
    def skipped(self) -> int:
        return sum(1 for r in self.results if r.status == TestStatus.SKIPPED)

    @property
    def errors(self) -> int:
        return sum(1 for r in self.results if r.status == TestStatus.ERROR)

    @property
    def total(self) -> int:
        return len(self.results)

    @property
    def all_passed(self) -> bool:
        return self.failed == 0 and self.errors == 0

    def to_dict(self) -> Dict[str, Any]:
        return {
            "prompt_name": self.prompt_name,
            "version": self.version,
            "results": [r.to_dict() for r in self.results],
            "summary": {
                "total": self.total,
                "passed": self.passed,
                "failed": self.failed,
                "skipped": self.skipped,
                "errors": self.errors,
            },
            "total_duration_ms": self.total_duration_ms,
            "all_passed": self.all_passed,
        }


class PromptTester:
    """Tests prompts against defined assertions."""

    # Built-in assertion types
    ASSERTION_TYPES = {
        "contains": "_assert_contains",
        "not_contains": "_assert_not_contains",
        "starts_with": "_assert_starts_with",
        "ends_with": "_assert_ends_with",
        "matches_regex": "_assert_matches_regex",
        "min_length": "_assert_min_length",
        "max_length": "_assert_max_length",
        "has_variable": "_assert_has_variable",
        "renders_without_error": "_assert_renders",
        "token_count_below": "_assert_token_count",
        "no_pii": "_assert_no_pii",
        "json_valid": "_assert_json_output",
        "word_count_between": "_assert_word_count",
    }

    def __init__(self):
        self._custom_assertions: Dict[str, Callable] = {}

    def register_assertion(self, name: str, fn: Callable) -> None:
        """Register a custom assertion function."""
        self._custom_assertions[name] = fn

    def test(self, prompt: Prompt, test_cases: List[TestCase], version: Optional[int] = None) -> TestReport:
        """Run all test cases against a prompt."""
        ver = prompt.get_version(version)
        if ver is None:
            ver = prompt.latest
        if ver is None:
            report = TestReport(prompt_name=prompt.name, version=0)
            report.results.append(TestResult(
                test_name="__setup__",
                status=TestStatus.ERROR,
                message="Prompt has no versions",
            ))
            return report

        report = TestReport(prompt_name=prompt.name, version=ver.version)
        start = time.monotonic()

        for tc in test_cases:
            result = self._run_test_case(ver, tc)
            report.results.append(result)

        report.total_duration_ms = (time.monotonic() - start) * 1000
        return report

    def test_content(self, content: str, test_cases: List[TestCase], name: str = "unnamed") -> TestReport:
        """Test raw prompt content."""
        ver = PromptVersion(version=1, content=content)
        prompt = Prompt(name=name)
        prompt.versions = [ver]
        return self.test(prompt, test_cases)

    def _run_test_case(self, version: PromptVersion, tc: TestCase) -> TestResult:
        """Run a single test case."""
        start = time.monotonic()

        # Render with variables if provided
        try:
            if tc.variables:
                rendered = version.render(tc.variables)
            else:
                rendered = version.content
        except ValueError as e:
            duration = (time.monotonic() - start) * 1000
            return TestResult(
                test_name=tc.name,
                status=TestStatus.ERROR,
                message=f"Render error: {e}",
                duration_ms=duration,
            )

        # Run each assertion
        failures = []
        for assertion in tc.assertions:
            atype = assertion.get("type", "")
            avalue = assertion.get("value")

            # Try built-in assertions
            method_name = self.ASSERTION_TYPES.get(atype)
            if method_name:
                method = getattr(self, method_name)
                ok, msg = method(version.content, rendered, avalue, tc.variables)
                if not ok:
                    failures.append(f"[{atype}] {msg}")
            # Try custom assertions
            elif atype in self._custom_assertions:
                try:
                    ok, msg = self._custom_assertions[atype](
                        version.content, rendered, avalue, tc.variables
                    )
                    if not ok:
                        failures.append(f"[{atype}] {msg}")
                except Exception as e:
                    failures.append(f"[{atype}] Error: {e}")
            else:
                failures.append(f"Unknown assertion type: {atype}")

        duration = (time.monotonic() - start) * 1000

        if failures:
            return TestResult(
                test_name=tc.name,
                status=TestStatus.FAILED,
                message="; ".join(failures),
                duration_ms=duration,
                details={"failures": failures},
            )

        return TestResult(
            test_name=tc.name,
            status=TestStatus.PASSED,
            message="All assertions passed",
            duration_ms=duration,
        )

    # ── Built-in assertions ──────────────────────────────────────

    def _assert_contains(self, raw: str, rendered: str, value: Any, _vars: Dict) -> tuple:
        if value and value in rendered:
            return True, ""
        return False, f"Expected '{value}' in prompt"

    def _assert_not_contains(self, raw: str, rendered: str, value: Any, _vars: Dict) -> tuple:
        if value and value not in rendered:
            return True, ""
        return False, f"Unexpected '{value}' found in prompt"

    def _assert_starts_with(self, raw: str, rendered: str, value: Any, _vars: Dict) -> tuple:
        if value and rendered.strip().startswith(value):
            return True, ""
        return False, f"Expected prompt to start with '{value}'"

    def _assert_ends_with(self, raw: str, rendered: str, value: Any, _vars: Dict) -> tuple:
        if value and rendered.strip().endswith(value):
            return True, ""
        return False, f"Expected prompt to end with '{value}'"

    def _assert_matches_regex(self, raw: str, rendered: str, value: Any, _vars: Dict) -> tuple:
        if value and re.search(value, rendered):
            return True, ""
        return False, f"Pattern '{value}' not found in prompt"

    def _assert_min_length(self, raw: str, rendered: str, value: Any, _vars: Dict) -> tuple:
        if isinstance(value, (int, float)) and len(rendered) >= value:
            return True, ""
        return False, f"Prompt length {len(rendered)} < minimum {value}"

    def _assert_max_length(self, raw: str, rendered: str, value: Any, _vars: Dict) -> tuple:
        if isinstance(value, (int, float)) and len(rendered) <= value:
            return True, ""
        return False, f"Prompt length {len(rendered)} > maximum {value}"

    def _assert_has_variable(self, raw: str, rendered: str, value: Any, _vars: Dict) -> tuple:
        pattern = r'\{\{\s*' + re.escape(str(value)) + r'\s*\}\}'
        if re.search(pattern, raw):
            return True, ""
        return False, f"Variable '{{{{ {value} }}}}' not found in prompt"

    def _assert_renders(self, raw: str, rendered: str, value: Any, variables: Dict) -> tuple:
        """Check that the prompt renders without Jinja2 errors."""
        try:
            from jinja2 import Template
            Template(raw).render(**(variables or {}))
            return True, ""
        except Exception as e:
            return False, f"Render error: {e}"

    def _assert_token_count(self, raw: str, rendered: str, value: Any, _vars: Dict) -> tuple:
        tokens = len(rendered) // 4  # rough estimate
        if isinstance(value, (int, float)) and tokens <= value:
            return True, ""
        return False, f"Estimated {tokens} tokens > limit {value}"

    def _assert_no_pii(self, raw: str, rendered: str, value: Any, _vars: Dict) -> tuple:
        """Check for common PII patterns."""
        pii_patterns = [
            (r'\b\d{3}-\d{2}-\d{4}\b', "SSN"),
            (r'\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Za-z]{2,}\b', "email"),
            (r'\b\d{4}[\s-]?\d{4}[\s-]?\d{4}[\s-]?\d{4}\b', "credit card"),
            (r'\b\d{3}[\s.-]?\d{3}[\s.-]?\d{4}\b', "phone number"),
        ]
        found = []
        for pattern, pii_type in pii_patterns:
            if re.search(pattern, rendered):
                found.append(pii_type)
        if found:
            return False, f"PII detected: {', '.join(found)}"
        return True, ""

    def _assert_json_output(self, raw: str, rendered: str, value: Any, _vars: Dict) -> tuple:
        """Check if the prompt asks for JSON output with proper format."""
        json_indicators = ["json", "JSON", '```json', '"type":', '{"']
        has_json = any(ind in rendered for ind in json_indicators)
        if has_json or (value and value == "optional"):
            return True, ""
        return False, "Prompt doesn't appear to specify JSON output format"

    def _assert_word_count(self, raw: str, rendered: str, value: Any, _vars: Dict) -> tuple:
        """Check word count is within range."""
        words = len(rendered.split())
        if isinstance(value, dict):
            min_words = value.get("min", 0)
            max_words = value.get("max", float("inf"))
            if min_words <= words <= max_words:
                return True, ""
            return False, f"Word count {words} not in range [{min_words}, {max_words}]"
        return True, ""


def load_test_cases(path: Path) -> List[TestCase]:
    """Load test cases from a JSON or YAML file."""
    if not path.exists():
        raise FileNotFoundError(f"Test file not found: {path}")

    content = path.read_text()

    if path.suffix in (".yaml", ".yml"):
        try:
            import yaml
            data = yaml.safe_load(content)
        except ImportError:
            raise ImportError("PyYAML required for YAML test files: pip install pyyaml")
    elif path.suffix == ".json":
        data = json.loads(content)
    else:
        raise ValueError(f"Unsupported test file format: {path.suffix}")

    if isinstance(data, dict):
        cases = data.get("tests", data.get("test_cases", []))
    elif isinstance(data, list):
        cases = data
    else:
        raise ValueError("Test file must contain a list or dict with 'tests' key")

    return [TestCase.from_dict(tc) for tc in cases]
