"""Property-Based Test (PBT) Generation from Replays.

Generates Hypothesis property-based tests from Hikaflow replay transcripts.
Unlike standard unit tests that use fixed inputs, PBTs verify invariants
across a wide range of generated inputs, catching edge cases.

Key features:
- Extracts invariants from recorded behavior
- Generates Hypothesis strategies from observed data
- Creates tests that verify properties, not just examples
- Catches edge cases that example-based tests miss

Research shows PBTs expose correctness gaps in 30-32% of solutions
that pass unit tests (arXiv:2506.18315).
"""

from __future__ import annotations

import json
import os
import re
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Set

try:
    import anthropic

    HAS_ANTHROPIC = True
except ImportError:
    HAS_ANTHROPIC = False


@dataclass
class ExtractedProperty:
    """A property extracted from replay behavior."""

    name: str
    description: str
    invariant: str  # Python expression that should always be True
    precondition: Optional[str] = None  # When the property applies
    confidence: float = 0.7


@dataclass
class GeneratedStrategy:
    """A Hypothesis strategy generated from observed data."""

    name: str
    code: str  # Hypothesis strategy code
    description: str
    source_field: str  # Which transcript field it came from


@dataclass
class PropertyBasedTest:
    """A generated property-based test."""

    name: str
    code: str
    strategies: List[str]  # Strategy names used
    properties: List[str]  # Property names tested
    imports: List[str]
    description: str
    confidence: float


@dataclass
class PBTGenerationResult:
    """Result of property-based test generation."""

    tests: List[PropertyBasedTest]
    strategies: List[GeneratedStrategy]
    properties: List[ExtractedProperty]
    conftest_code: str
    errors: List[str] = field(default_factory=list)


# Prompt for property extraction
PROPERTY_EXTRACTION_PROMPT = """You are an expert at property-based testing. Analyze this recorded I/O transcript and extract testable properties (invariants).

## Recorded Behavior

### HTTP Calls
{http_calls}

### SQL Queries
{sql_queries}

### Function Inputs/Outputs
Test: {test_name}
Error: {error_type}: {error_message}

## Instructions

Extract properties that should ALWAYS be true about this code's behavior. Focus on:
1. Response format invariants (e.g., "response always has 'id' field")
2. Data type invariants (e.g., "user_id is always positive integer")
3. Relationship invariants (e.g., "output list length <= input length")
4. State invariants (e.g., "balance never goes negative")
5. Ordering invariants (e.g., "results are sorted by date")

For each property, specify:
- name: descriptive snake_case name
- description: human-readable explanation
- invariant: Python boolean expression
- precondition: when this property applies (or null)
- confidence: 0.0-1.0

Respond with JSON array:
```json
[
  {{
    "name": "response_has_id",
    "description": "API response always contains an id field",
    "invariant": "'id' in response",
    "precondition": "response.status_code == 200",
    "confidence": 0.9
  }}
]
```
"""

# Prompt for strategy generation
STRATEGY_GENERATION_PROMPT = """Generate Hypothesis strategies based on observed data patterns.

## Observed Data

### HTTP Request Bodies
{request_bodies}

### SQL Query Parameters
{sql_params}

### Response Structures
{response_structures}

## Instructions

Generate Hypothesis strategies that can produce data similar to what was observed.
Each strategy should:
1. Match the observed data types and structures
2. Include edge cases (empty strings, zero, negative, unicode, etc.)
3. Be composable with other strategies

Respond with JSON array:
```json
[
  {{
    "name": "user_ids",
    "code": "st.integers(min_value=1, max_value=1000000)",
    "description": "Valid user IDs based on observed range",
    "source_field": "request.user_id"
  }},
  {{
    "name": "email_addresses",
    "code": "st.emails()",
    "description": "Email addresses",
    "source_field": "request.email"
  }}
]
```

Use these Hypothesis strategies:
- st.integers(), st.floats(), st.text(), st.binary()
- st.lists(), st.dictionaries(), st.tuples()
- st.emails(), st.uuids(), st.datetimes()
- st.one_of(), st.sampled_from()
- st.builds() for complex objects
"""

# Prompt for test generation
PBT_GENERATION_PROMPT = """Generate a Hypothesis property-based test that verifies these properties.

## Properties to Test
{properties}

## Available Strategies
{strategies}

## Code Context
Function: {function_name}
File: {file_path}
Error that occurred: {error_type}: {error_message}

## Instructions

Generate a @given decorated test function that:
1. Uses the provided strategies to generate inputs
2. Tests all the specified properties
3. Includes @settings for reproducibility
4. Has clear assertion messages

Example format:
```python
@settings(max_examples=100, deadline=None)
@given(
    user_id=user_ids(),
    amount=amounts(),
)
def test_transfer_properties(user_id, amount):
    \"\"\"Property: transfers maintain balance invariants.\"\"\"
    # Setup
    initial_balance = get_balance(user_id)

    # Action
    result = transfer(user_id, amount)

    # Properties
    assert result.status in ['success', 'insufficient_funds'], "Invalid status"
    if result.status == 'success':
        assert get_balance(user_id) == initial_balance - amount, "Balance mismatch"
```

Respond with JSON:
```json
{{
  "name": "test_<descriptive_name>_properties",
  "code": "...",
  "strategies_used": ["user_ids", "amounts"],
  "properties_tested": ["balance_invariant", "status_valid"],
  "imports": ["from hypothesis import given, settings, strategies as st"],
  "description": "Property-based test for transfer invariants",
  "confidence": 0.8
}}
```
"""


def _format_http_calls(transcript: Dict[str, Any]) -> str:
    """Format HTTP calls for prompt."""
    http_calls = transcript.get("http", [])
    if not http_calls:
        return "No HTTP calls recorded."

    lines = []
    for i, call in enumerate(http_calls[:5], 1):
        method = call.get("method", "GET")
        url = call.get("url", "?")
        status = call.get("status", "?")
        req_body = call.get("request_body", {})
        resp_body = call.get("response_body", {})

        lines.append(f"{i}. {method} {url} -> {status}")
        if req_body:
            lines.append(f"   Request: {json.dumps(req_body)[:200]}")
        if resp_body:
            lines.append(f"   Response: {json.dumps(resp_body)[:200]}")

    return "\n".join(lines)


def _format_sql_queries(transcript: Dict[str, Any]) -> str:
    """Format SQL queries for prompt."""
    sql_queries = transcript.get("sql", [])
    if not sql_queries:
        return "No SQL queries recorded."

    lines = []
    for i, q in enumerate(sql_queries[:5], 1):
        query = q.get("query", "?")[:150]
        params = q.get("params", [])
        lines.append(f"{i}. {query}")
        if params:
            lines.append(f"   Params: {params}")

    return "\n".join(lines)


def _extract_request_bodies(transcript: Dict[str, Any]) -> str:
    """Extract request body structures from HTTP calls."""
    http_calls = transcript.get("http", [])
    bodies = []

    for call in http_calls[:5]:
        body = call.get("request_body")
        if body:
            bodies.append(json.dumps(body, indent=2)[:300])

    return "\n---\n".join(bodies) if bodies else "No request bodies."


def _extract_response_structures(transcript: Dict[str, Any]) -> str:
    """Extract response structures from HTTP calls."""
    http_calls = transcript.get("http", [])
    structures = []

    for call in http_calls[:5]:
        body = call.get("response_body")
        if body:
            structures.append(json.dumps(body, indent=2)[:300])

    return "\n---\n".join(structures) if structures else "No response bodies."


def _extract_sql_params(transcript: Dict[str, Any]) -> str:
    """Extract SQL parameters."""
    sql_queries = transcript.get("sql", [])
    params = []

    for q in sql_queries[:5]:
        p = q.get("params", [])
        if p:
            params.append(str(p))

    return "\n".join(params) if params else "No SQL parameters."


def _call_llm(prompt: str) -> str:
    """Call LLM for generation."""
    # Prefer DeepSeek/OpenAI-compatible path via central config
    try:
        from autodebug.llm_config import get_llm_config, create_openai_client
        config = get_llm_config()
        if config.is_openai_compatible:
            client = create_openai_client(config)
            response = client.chat.completions.create(
                model=config.default_model,
                max_tokens=3000,
                messages=[{"role": "user", "content": prompt}],
            )
            return response.choices[0].message.content
    except (ImportError, ValueError):
        pass

    if not HAS_ANTHROPIC or not os.environ.get("ANTHROPIC_API_KEY"):
        raise RuntimeError("No LLM API key available")

    client = anthropic.Anthropic()
    response = client.messages.create(
        model="claude-sonnet-4-20250514",
        max_tokens=3000,
        messages=[{"role": "user", "content": prompt}],
    )
    return response.content[0].text


def _parse_json_response(response: str) -> Any:
    """Parse JSON from LLM response."""
    # Try to extract JSON from code blocks
    json_match = re.search(r"```(?:json)?\s*([\[\{].*?[\]\}])\s*```", response, re.DOTALL)
    if json_match:
        try:
            return json.loads(json_match.group(1))
        except json.JSONDecodeError:
            pass

    # Try to find raw JSON
    for pattern in [r"\[.*\]", r"\{.*\}"]:
        match = re.search(pattern, response, re.DOTALL)
        if match:
            try:
                return json.loads(match.group())
            except json.JSONDecodeError:
                continue

    return None


def extract_properties(
    transcript: Dict[str, Any],
    context: Dict[str, Any],
) -> List[ExtractedProperty]:
    """Extract testable properties from a replay transcript.

    Args:
        transcript: I/O transcript with http, sql, etc.
        context: Test context with error info

    Returns:
        List of extracted properties
    """
    prompt = PROPERTY_EXTRACTION_PROMPT.format(
        http_calls=_format_http_calls(transcript),
        sql_queries=_format_sql_queries(transcript),
        test_name=context.get("test_name", "unknown"),
        error_type=context.get("error_type", "Exception"),
        error_message=context.get("error_message", "")[:300],
    )

    try:
        response = _call_llm(prompt)
        parsed = _parse_json_response(response)

        if not parsed or not isinstance(parsed, list):
            return []

        return [
            ExtractedProperty(
                name=p.get("name", f"property_{i}"),
                description=p.get("description", ""),
                invariant=p.get("invariant", "True"),
                precondition=p.get("precondition"),
                confidence=p.get("confidence", 0.7),
            )
            for i, p in enumerate(parsed)
        ]
    except Exception:
        return []


def generate_strategies(transcript: Dict[str, Any]) -> List[GeneratedStrategy]:
    """Generate Hypothesis strategies from observed data.

    Args:
        transcript: I/O transcript

    Returns:
        List of generated strategies
    """
    prompt = STRATEGY_GENERATION_PROMPT.format(
        request_bodies=_extract_request_bodies(transcript),
        sql_params=_extract_sql_params(transcript),
        response_structures=_extract_response_structures(transcript),
    )

    try:
        response = _call_llm(prompt)
        parsed = _parse_json_response(response)

        if not parsed or not isinstance(parsed, list):
            return _generate_fallback_strategies(transcript)

        import ast as _ast

        result = []
        for i, s in enumerate(parsed):
            code = s.get("code", "st.none()")
            # Validate strategy code: must parse and not contain dangerous patterns
            try:
                _ast.parse(code, mode="eval")
            except SyntaxError:
                code = "st.none()"
            result.append(GeneratedStrategy(
                name=s.get("name", f"strategy_{i}"),
                code=code,
                description=s.get("description", ""),
                source_field=s.get("source_field", "unknown"),
            ))
        return result
    except Exception:
        return _generate_fallback_strategies(transcript)


def _generate_fallback_strategies(transcript: Dict[str, Any]) -> List[GeneratedStrategy]:
    """Generate basic strategies without LLM."""
    strategies = []

    # Analyze HTTP bodies for types
    http_calls = transcript.get("http", [])
    observed_types: Dict[str, Set[str]] = {}

    for call in http_calls:
        for body in [call.get("request_body"), call.get("response_body")]:
            if isinstance(body, dict):
                for key, value in body.items():
                    if key not in observed_types:
                        observed_types[key] = set()
                    observed_types[key].add(type(value).__name__)

    # Generate strategies based on observed types
    for field_name, types in observed_types.items():
        if "int" in types:
            strategies.append(
                GeneratedStrategy(
                    name=f"{field_name}_values",
                    code="st.integers()",
                    description=f"Integer values for {field_name}",
                    source_field=field_name,
                )
            )
        elif "str" in types:
            strategies.append(
                GeneratedStrategy(
                    name=f"{field_name}_values",
                    code="st.text(min_size=1, max_size=100)",
                    description=f"String values for {field_name}",
                    source_field=field_name,
                )
            )
        elif "list" in types:
            strategies.append(
                GeneratedStrategy(
                    name=f"{field_name}_values",
                    code="st.lists(st.text(), max_size=10)",
                    description=f"List values for {field_name}",
                    source_field=field_name,
                )
            )

    return strategies


def generate_pbt(
    properties: List[ExtractedProperty],
    strategies: List[GeneratedStrategy],
    context: Dict[str, Any],
) -> Optional[PropertyBasedTest]:
    """Generate a property-based test.

    Args:
        properties: Properties to test
        strategies: Available strategies
        context: Test context

    Returns:
        Generated PBT or None
    """
    if not properties:
        return None

    # Format properties for prompt
    props_str = "\n".join(
        f"- {p.name}: {p.description}\n  Invariant: {p.invariant}"
        for p in properties
    )

    # Format strategies for prompt
    strats_str = "\n".join(
        f"- {s.name}(): {s.code}\n  {s.description}"
        for s in strategies
    )

    prompt = PBT_GENERATION_PROMPT.format(
        properties=props_str,
        strategies=strats_str,
        function_name=context.get("test_name", "unknown"),
        file_path=context.get("file", "unknown"),
        error_type=context.get("error_type", "Exception"),
        error_message=context.get("error_message", "")[:200],
    )

    try:
        response = _call_llm(prompt)
        parsed = _parse_json_response(response)

        if not parsed:
            return _generate_fallback_pbt(properties, strategies, context)

        code = parsed.get("code", "")

        # Validate LLM-generated code: syntax check + dangerous pattern blocklist
        if code:
            import ast as _ast

            try:
                _ast.parse(code)
            except SyntaxError:
                return _generate_fallback_pbt(properties, strategies, context)

            _DANGEROUS_PATTERNS = ["eval(", "exec(", "__import__(", "subprocess", "os.system(", "shutil.rmtree("]
            code_lower = code.lower()
            if any(pat.lower() in code_lower for pat in _DANGEROUS_PATTERNS):
                return _generate_fallback_pbt(properties, strategies, context)

        return PropertyBasedTest(
            name=parsed.get("name", "test_properties"),
            code=code,
            strategies=parsed.get("strategies_used", []),
            properties=parsed.get("properties_tested", []),
            imports=parsed.get("imports", ["from hypothesis import given, settings, strategies as st"]),
            description=parsed.get("description", ""),
            confidence=parsed.get("confidence", 0.7),
        )
    except Exception:
        return _generate_fallback_pbt(properties, strategies, context)


def _generate_fallback_pbt(
    properties: List[ExtractedProperty],
    strategies: List[GeneratedStrategy],
    context: Dict[str, Any],
) -> PropertyBasedTest:
    """Generate a basic PBT without LLM."""
    test_name = context.get("test_name", "unknown").replace("test_", "")

    # Build strategy parameters
    strat_params = []
    for s in strategies[:3]:  # Limit to 3 strategies
        strat_params.append(f"    {s.name}={s.code},")

    strat_code = "\n".join(strat_params) if strat_params else "    value=st.integers(),"

    # Build property assertions
    assertions = []
    for p in properties[:5]:  # Limit to 5 properties
        assertions.append(f"    # Property: {p.description}")
        if p.precondition:
            assertions.append(f"    if {p.precondition}:")
            assertions.append(f"        assert {p.invariant}, '{p.name} violated'")
        else:
            assertions.append(f"    assert {p.invariant}, '{p.name} violated'")

    assertions_code = "\n".join(assertions) if assertions else "    pass  # Add assertions"

    code = f'''@settings(max_examples=100, deadline=None)
@given(
{strat_code}
)
def test_{test_name}_properties(**kwargs):
    """Property-based test generated from Hikaflow replay.

    Tests these properties:
    {chr(10).join(f"- {p.name}: {p.description}" for p in properties[:5])}
    """
    # Call the function under test with kwargs.
    # Replace 'function_under_test' with the actual target function.
    pytest.skip("Generated PBT: replace function_under_test() with actual call")
    # result = function_under_test(**kwargs)

{assertions_code}
'''

    return PropertyBasedTest(
        name=f"test_{test_name}_properties",
        code=code,
        strategies=[s.name for s in strategies[:3]],
        properties=[p.name for p in properties[:5]],
        imports=[
            "from hypothesis import given, settings, strategies as st",
            "import pytest",
        ],
        description=f"Property-based test for {test_name}",
        confidence=0.5,
    )


def generate_pbt_from_replay(
    run_id: str,
    transcript: Dict[str, Any],
    context: Dict[str, Any],
) -> PBTGenerationResult:
    """Generate property-based tests from a replay.

    Args:
        run_id: Hikaflow run ID
        transcript: I/O transcript
        context: Test context

    Returns:
        PBTGenerationResult with tests, strategies, and properties
    """
    errors = []

    # Step 1: Extract properties
    try:
        properties = extract_properties(transcript, context)
    except Exception as e:
        errors.append(f"Property extraction failed: {e}")
        properties = []

    # Step 2: Generate strategies
    try:
        strategies = generate_strategies(transcript)
    except Exception as e:
        errors.append(f"Strategy generation failed: {e}")
        strategies = []

    # Step 3: Generate PBT
    tests = []
    if properties or strategies:
        try:
            pbt = generate_pbt(properties, strategies, context)
            if pbt:
                tests.append(pbt)
        except Exception as e:
            errors.append(f"PBT generation failed: {e}")

    # Generate conftest code
    conftest_code = _generate_conftest(strategies)

    return PBTGenerationResult(
        tests=tests,
        strategies=strategies,
        properties=properties,
        conftest_code=conftest_code,
        errors=errors,
    )


def _generate_conftest(strategies: List[GeneratedStrategy]) -> str:
    """Generate conftest.py code with strategies as fixtures."""
    lines = [
        '"""Hypothesis strategies generated from Hikaflow replays."""',
        "",
        "from hypothesis import strategies as st",
        "import pytest",
        "",
        "# Generated strategies",
        "",
    ]

    for s in strategies:
        lines.extend([
            "@st.composite",
            f"def {s.name}(draw):",
            f'    """Strategy for {s.source_field}.',
            "",
            f"    {s.description}",
            '    """',
            f"    return draw({s.code})",
            "",
        ])

    return "\n".join(lines)


def format_pbt_file(result: PBTGenerationResult) -> str:
    """Format PBT result as a complete test file.

    Args:
        result: PBTGenerationResult

    Returns:
        Complete test file content
    """
    # Collect imports
    all_imports = {
        "from hypothesis import given, settings, strategies as st",
        "import pytest",
    }

    for test in result.tests:
        for imp in test.imports:
            all_imports.add(imp)

    lines = [
        '"""Property-based tests generated from Hikaflow replay data.',
        "",
        "These tests verify invariants across a wide range of inputs,",
        "catching edge cases that example-based tests might miss.",
        '"""',
        "",
        *sorted(all_imports),
        "",
        "",
    ]

    # Add strategy definitions
    if result.strategies:
        lines.append("# Custom strategies from observed data")
        for s in result.strategies:
            lines.append(f"def {s.name}():")
            lines.append(f'    """{s.description}"""')
            lines.append(f"    return {s.code}")
            lines.append("")
        lines.append("")

    # Add tests
    for test in result.tests:
        lines.append(f"# {test.description}")
        lines.append(f"# Confidence: {test.confidence:.0%}")
        lines.append(test.code)
        lines.append("")
        lines.append("")

    return "\n".join(lines)
