"""
SOQL Rewriter: Enforce row limits and add ORDER BY for safety.
Rewrites queries to protect free tier and ensure stable pagination.
"""

import re
from typing import Tuple, Optional


def enforce_row_limit(
    soql: str, max_rows: int, tier: str = "free"
) -> Tuple[str, bool, Optional[str]]:
    """
    Enforce row limit on SOQL query.

    Args:
        soql: Original SOQL query
        max_rows: Maximum rows allowed for this tier
        tier: User's tier (for logging/messaging)

    Returns:
        Tuple of (rewritten_soql, was_modified, modification_message)
    """
    soql_upper = soql.upper()
    was_modified = False
    message = None

    # Extract current LIMIT if exists
    limit_match = re.search(r"\bLIMIT\s+(\d+)", soql_upper)
    current_limit = int(limit_match.group(1)) if limit_match else None

    # Check ORDER BY presence
    has_order_by = "ORDER BY" in soql_upper

    modified_soql = soql

    # 1. Enforce LIMIT
    if current_limit is None:
        # No LIMIT - add one
        modified_soql = f"{modified_soql.rstrip()} LIMIT {max_rows}"
        was_modified = True
        message = f"Added LIMIT {max_rows} (no limit specified, capped for {tier} tier)"

    elif current_limit > max_rows:
        # LIMIT exceeds tier max - reduce it
        modified_soql = re.sub(
            r"\bLIMIT\s+\d+", f"LIMIT {max_rows}", modified_soql, flags=re.IGNORECASE
        )
        was_modified = True
        message = f"Reduced LIMIT from {current_limit} to {max_rows} ({tier} tier max)"

    # 2. Add ORDER BY if missing (for pagination stability) whenever LIMIT exists.
    has_limit_after = bool(re.search(r"\bLIMIT\s+\d+", modified_soql, re.IGNORECASE))
    if has_limit_after and not has_order_by:
        limit_match_after = re.search(r"\bLIMIT\s+\d+", modified_soql, re.IGNORECASE)
        insert_pos = limit_match_after.start() if limit_match_after else len(modified_soql)
        modified_soql = (
            modified_soql[:insert_pos].rstrip()
            + " ORDER BY Id "
            + modified_soql[insert_pos:]
        )
        was_modified = True
        if message:
            message += "; Added ORDER BY Id for stable pagination"
        else:
            message = "Added ORDER BY Id for stable pagination"

    return modified_soql, was_modified, message


def estimate_query_complexity(soql: str) -> dict:
    """
    Estimate query complexity for COUNT() feasibility.

    Returns:
        {
            'is_simple': bool,  # Can we run COUNT()?
            'has_groupby': bool,
            'has_subquery': bool,
            'has_aggregate': bool,
            'object_name': str or None
        }
    """
    soql_upper = soql.upper()

    # Extract object name
    from_match = re.search(r"\bFROM\s+(\w+)", soql_upper)
    object_name = from_match.group(1) if from_match else None

    # Check for complex features
    has_groupby = "GROUP BY" in soql_upper
    has_subquery = "(" in soql and "SELECT" in soql_upper[soql_upper.find("(") :]
    has_aggregate = any(
        agg in soql_upper for agg in ["COUNT(", "SUM(", "AVG(", "MIN(", "MAX("]
    )

    # Simple query = SELECT fields FROM object WHERE ... (no GROUP BY, no subquery, no aggregates)
    is_simple = not (has_groupby or has_subquery or has_aggregate)

    return {
        "is_simple": is_simple,
        "has_groupby": has_groupby,
        "has_subquery": has_subquery,
        "has_aggregate": has_aggregate,
        "object_name": object_name,
    }


def generate_count_query(soql: str) -> Optional[str]:
    """
    Generate a COUNT() query from the original SOQL (only for simple queries).

    Args:
        soql: Original SOQL query

    Returns:
        COUNT query string or None if not feasible
    """
    complexity = estimate_query_complexity(soql)

    if not complexity["is_simple"]:
        return None

    # Extract FROM and WHERE clauses
    from_match = re.search(r"\bFROM\s+\w+", soql, re.IGNORECASE)
    if not from_match:
        return None

    from_clause = from_match.group(0)

    # Extract WHERE clause if exists
    where_match = re.search(r"\bWHERE\b.*", soql, re.IGNORECASE)
    where_clause = ""
    if where_match:
        where_clause = where_match.group(0)
        # Remove ORDER BY, LIMIT from WHERE clause
        where_clause = re.sub(
            r"\b(ORDER BY|LIMIT)\b.*", "", where_clause, flags=re.IGNORECASE
        ).strip()

    count_query = f"SELECT COUNT() {from_clause}"
    if where_clause:
        count_query += f" {where_clause}"

    return count_query


def add_suggested_fixes(soql: str, tier: str = "free", max_rows: int = 15000) -> list:
    """
    Generate suggested fixes for SOQL query.

    Returns:
        List of {fix_type, description, suggested_soql}
    """
    fixes = []
    soql_upper = soql.upper()

    # Check for LIMIT
    limit_match = re.search(r"\bLIMIT\s+(\d+)", soql_upper)
    if not limit_match:
        fixed_soql = f"{soql.rstrip()} LIMIT {max_rows}"
        fixes.append(
            {
                "fix_type": "add_limit",
                "description": f"Add LIMIT {max_rows} to cap results",
                "suggested_soql": fixed_soql,
                "reason": "Prevents large result sets and excessive API calls",
            }
        )
    elif int(limit_match.group(1)) > max_rows:
        fixed_soql = re.sub(
            r"\bLIMIT\s+\d+", f"LIMIT {max_rows}", soql, flags=re.IGNORECASE
        )
        fixes.append(
            {
                "fix_type": "reduce_limit",
                "description": f"Reduce LIMIT to {max_rows} ({tier} tier maximum)",
                "suggested_soql": fixed_soql,
                "reason": f"{tier.title()} tier is limited to {max_rows} rows per query",
            }
        )

    # Check for ORDER BY
    if "LIMIT" in soql_upper and "ORDER BY" not in soql_upper:
        # Find position before LIMIT
        limit_pos = soql_upper.find("LIMIT")
        fixed_soql = soql[:limit_pos].rstrip() + " ORDER BY Id " + soql[limit_pos:]
        fixes.append(
            {
                "fix_type": "add_order_by",
                "description": "Add ORDER BY Id for stable pagination",
                "suggested_soql": fixed_soql,
                "reason": "Ensures consistent results across pages",
            }
        )

    # Check for large objects without WHERE
    large_objects = ["ACCOUNT", "CONTACT", "LEAD", "OPPORTUNITY", "CASE"]
    for obj in large_objects:
        if f"FROM {obj}" in soql_upper and "WHERE" not in soql_upper:
            # Suggest adding WHERE with recent records
            upper = soql_upper
            clause_positions = [
                upper.find(" GROUP BY "),
                upper.find(" ORDER BY "),
                upper.find(" LIMIT "),
                upper.find(" OFFSET "),
            ]
            clause_positions = [pos for pos in clause_positions if pos >= 0]
            insert_pos = min(clause_positions) if clause_positions else len(soql)
            fixed_soql = (
                soql[:insert_pos].rstrip()
                + " WHERE CreatedDate >= LAST_N_DAYS:30 "
                + soql[insert_pos:].lstrip()
            ).strip()
            fixes.append(
                {
                    "fix_type": "add_where_clause",
                    "description": f"Add WHERE clause to filter {obj} records",
                    "suggested_soql": fixed_soql,
                    "reason": f"{obj} is a large object. Filter by date to reduce result size.",
                }
            )
            break

    return fixes


# Example usage in endpoint
def validate_and_rewrite_soql(soql: str, tier: str, max_rows: int) -> dict:
    """
    Complete validation and rewrite pipeline.

    Returns:
        {
            'original_soql': str,
            'rewritten_soql': str,
            'was_modified': bool,
            'modification_message': str,
            'suggested_fixes': list,
            'complexity': dict,
            'count_query': str or None
        }
    """
    # 1. Enforce limits
    rewritten_soql, was_modified, mod_message = enforce_row_limit(soql, max_rows, tier)

    # 2. Analyze complexity
    complexity = estimate_query_complexity(soql)

    # 3. Generate COUNT query if simple
    count_query = generate_count_query(soql) if complexity["is_simple"] else None

    # 4. Generate suggested fixes
    suggested_fixes = add_suggested_fixes(soql, tier, max_rows)

    return {
        "original_soql": soql,
        "rewritten_soql": rewritten_soql,
        "was_modified": was_modified,
        "modification_message": mod_message,
        "suggested_fixes": suggested_fixes,
        "complexity": complexity,
        "count_query": count_query,
    }
