# extractor_helpers.py
import os
import json
import urllib.parse
from typing import Dict, List, Optional, Any
from fastmcp import Context  # ✅ FastMCP 2.x import

from utils.config import load_config, load_jmeter_config

# === Global configuration ===
CONFIG = load_config()
ARTIFACTS_PATH = CONFIG["artifacts"]["artifacts_path"]
JMETER_CONFIG = load_jmeter_config()

# Import necessary modules for JMeter JMX generation
import xml.etree.ElementTree as ET  # Needed for creating empty hashTree elements

# Import necessary modules for JMeter JMX generation
from services.jmx.post_processor import (
    create_boundary_extractor,
    create_json_extractor,
    create_regex_extractor
)

from services.helpers.hostname_helpers import _categorize_hostname  # For OAuth extractor filtering

# ============================================================
# Helper Functions - Correlation & Extractor Support
# ============================================================

def _load_correlation_naming(test_run_id: str) -> Optional[Dict[str, Any]]:
    """
    Load correlation_naming.json if it exists.
    
    This file is generated by applying Cursor Rules to correlation_spec.json
    and contains JMeter variable names and extractor configurations.
    
    Args:
        test_run_id: The test run identifier
        
    Returns:
        The correlation naming data, or None if file doesn't exist
    """
    naming_path = os.path.join(ARTIFACTS_PATH, test_run_id, "jmeter", "correlation_naming.json")
    if os.path.isfile(naming_path):
        with open(naming_path, "r", encoding="utf-8") as f:
            return json.load(f)
    return None


def _load_correlation_spec(test_run_id: str) -> Optional[Dict[str, Any]]:
    """
    Load correlation_spec.json if it exists.
    
    This file contains the detailed correlation analysis including actual values
    and usage locations needed for variable substitution.
    
    Args:
        test_run_id: The test run identifier
        
    Returns:
        The correlation spec data, or None if file doesn't exist
    """
    spec_path = os.path.join(ARTIFACTS_PATH, test_run_id, "jmeter", "correlation_spec.json")
    if os.path.isfile(spec_path):
        with open(spec_path, "r", encoding="utf-8") as f:
            return json.load(f)
    return None


def _normalize_url(url: str) -> str:
    """
    Normalize a URL for comparison by removing query parameters and fragments.
    
    Args:
        url: The URL to normalize
        
    Returns:
        The normalized URL (scheme://host/path)
    """
    parsed = urllib.parse.urlparse(url)
    return f"{parsed.scheme}://{parsed.netloc}{parsed.path}"


def _build_extractor_map(correlation_naming: Dict[str, Any]) -> Dict[str, List[Dict]]:
    """
    Build a mapping from source URL to list of extractor configurations.
    
    This allows quick lookup when creating HTTP samplers to determine
    which extractors should be added to each sampler's hashTree.
    
    Args:
        correlation_naming: The loaded correlation_naming.json data
        
    Returns:
        Dictionary mapping normalized URLs to lists of variable configs
    """
    extractor_map = {}
    
    for var in correlation_naming.get("variables", []):
        source_url = var.get("source_request_url", "")
        if not source_url:
            continue
            
        # Store both exact URL and normalized version for flexible matching
        normalized = _normalize_url(source_url)
        
        if normalized not in extractor_map:
            extractor_map[normalized] = []
        extractor_map[normalized].append(var)
        
        # Also store exact URL if different from normalized
        if source_url != normalized:
            if source_url not in extractor_map:
                extractor_map[source_url] = []
            # Avoid duplicates
            if var not in extractor_map[source_url]:
                extractor_map[source_url].append(var)
    
    return extractor_map


def _find_extractors_for_url(url: str, extractor_map: Dict[str, List[Dict]]) -> List[Dict]:
    """
    Find extractor configurations that match a given URL.
    
    Attempts multiple matching strategies:
    1. Exact URL match
    2. Normalized URL match (without query params)
    3. Partial match (URL contains or is contained by source URL path)
    
    Args:
        url: The URL of the current HTTP sampler
        extractor_map: The pre-built extractor mapping
        
    Returns:
        List of variable configurations that should extract from this URL
    """
    # Try exact match first
    if url in extractor_map:
        return extractor_map[url]
    
    # Try normalized URL match
    normalized = _normalize_url(url)
    if normalized in extractor_map:
        return extractor_map[normalized]
    
    # Try partial path matching for edge cases
    parsed = urllib.parse.urlparse(url)
    url_path = parsed.path.rstrip('/')
    
    for source_url, extractors in extractor_map.items():
        source_parsed = urllib.parse.urlparse(source_url)
        source_path = source_parsed.path.rstrip('/')
        
        # Match if same host and path matches
        if (parsed.netloc == source_parsed.netloc and 
            (url_path == source_path or 
             url_path.startswith(source_path + '/') or
             source_path.startswith(url_path + '/'))):
            return extractors
    
    return []


def _infer_field_to_check(expression: str, extractor_type: str) -> str:
    """
    Infer the 'field_to_check' parameter for regex extractors.
    
    Based on the expression pattern, determine where JMeter should
    search for the value (body, headers, or URL).
    
    Args:
        expression: The regex or JSONPath expression
        extractor_type: "json_extractor" or "regex_extractor"
        
    Returns:
        One of: "body", "headers", "url"
    """
    if extractor_type == "json_extractor":
        return "body"  # JSON is always in response body
    
    expression_lower = expression.lower()
    
    # URL query parameter pattern (e.g., [?&]code=([^&]+))
    if "[?&]" in expression or "\\?" in expression:
        return "url"
    
    # Header patterns (e.g., X-Header: (.+), Set-Cookie:, Location:)
    header_indicators = [
        "set-cookie", "location:", "x-", "authorization:",
        "content-type:", "www-authenticate:", "bearer"
    ]
    for indicator in header_indicators:
        if indicator in expression_lower:
            return "headers"
    
    # Cookie extraction patterns
    if "cookie" in expression_lower:
        return "headers"
    
    # Default to body
    return "body"


def _create_extractor_element(var_config: Dict) -> Optional[ET.Element]:
    """
    Create the appropriate JMeter extractor element for a correlation variable.
    
    Based on the jmeter_extractor_type, creates either a JSON Extractor
    or Regular Expression Extractor with the configured parameters.
    
    Args:
        var_config: Variable configuration from correlation_naming.json
        
    Returns:
        The extractor XML element, or None if invalid config
    """
    extractor_type = var_config.get("jmeter_extractor_type", "")
    variable_name = var_config.get("variable_name", "")
    expression = var_config.get("jmeter_extractor_expression", "")
    testname = var_config.get("jmeter_extractor_name", f"Extract {variable_name}")
    
    if extractor_type == "boundary_extractor":
        left_boundary = var_config.get("left_boundary", "")
        right_boundary = var_config.get("right_boundary", "")
        if not variable_name or (not left_boundary and not right_boundary):
            return None
        return create_boundary_extractor(
            variable_name=variable_name,
            left_boundary=left_boundary,
            right_boundary=right_boundary,
            testname=testname
        )

    if not variable_name or not expression:
        return None
    
    if extractor_type == "json_extractor":
        return create_json_extractor(
            variable_name=variable_name,
            json_path=expression,
            testname=testname
        )
    elif extractor_type == "regex_extractor":
        field_to_check = _infer_field_to_check(expression, extractor_type)
        return create_regex_extractor(
            variable_name=variable_name,
            regex=expression,
            field_to_check=field_to_check,
            testname=testname
        )
    
    return None


def _should_add_oauth_extractor(
    hostname_category: str,
    url_path: str
) -> bool:
    """
    Determine if OAuth extractors should be added for this request.
    
    OAuth extractors should only be added for:
    1. Auth-related hostnames (auth, auth_internal) - where OAuth params originate
    2. Refresh token endpoints - regardless of hostname, if path contains both "refresh" and "token"
    
    Args:
        hostname_category: The category of the hostname ("auth", "auth_internal", "app", "static")
        url_path: The URL path to check for refresh token endpoints
        
    Returns:
        True if OAuth extractors should be added, False otherwise
    """
    # Always add for auth-related hostnames
    if hostname_category in ("auth", "auth_internal"):
        return True
    
    # Check for refresh token endpoint (case-insensitive)
    path_lower = url_path.lower()
    if "refresh" in path_lower and "token" in path_lower:
        return True
    
    # Skip for app/static hostnames (OAuth values are just echoed, not generated)
    return False


def _get_extractors_for_entry(
    entry: Dict, 
    extractor_map: Dict[str, List[Dict]],
    extracted_variables: Optional[set] = None,
    extractor_config: Optional[Dict] = None,
    hostname_patterns_config: Optional[Dict] = None
) -> List[ET.Element]:
    """
    Get all extractor elements that should be added to a sampler.
    
    Supports "first_occurrence" mode where extractors are only added the first
    time a variable is encountered, with exceptions for OAuth parameters.
    
    OAuth extractors are filtered to only apply on auth-related hostnames
    or refresh token endpoints.
    
    Args:
        entry: The network capture entry (contains URL)
        extractor_map: The pre-built extractor mapping
        extracted_variables: Set tracking which variables already have extractors (modified in-place)
        extractor_config: Configuration from jmeter_config.yaml for extractor placement
        hostname_patterns_config: Configuration for hostname categorization (for OAuth filtering)
        
    Returns:
        List of extractor XML elements to add to the sampler's hashTree
    """
    url = entry.get("url", "")
    if not url:
        return []
    
    var_configs = _find_extractors_for_url(url, extractor_map)
    extractors = []
    
    # Get extractor placement settings
    mode = "all_occurrences"  # Default to adding all extractors
    always_all = []  # OAuth variables that need special handling
    
    if extractor_config:
        mode = extractor_config.get("mode", "all_occurrences")
        always_all = extractor_config.get("always_all_occurrences", [])
    
    # Parse URL to get hostname and path for OAuth filtering
    parsed_url = urllib.parse.urlparse(url)
    hostname = parsed_url.netloc
    url_path = parsed_url.path
    
    # Determine hostname category for OAuth filtering
    hostname_category = "app"  # Default
    if hostname_patterns_config and hostname:
        hostname_category = _categorize_hostname(hostname, hostname_patterns_config)
    
    # Check once if this request qualifies for OAuth extractors
    oauth_allowed = _should_add_oauth_extractor(hostname_category, url_path)
    
    for var_config in var_configs:
        variable_name = var_config.get("variable_name", "")
        
        # Check if this is an OAuth variable
        is_oauth_variable = variable_name in always_all
        
        # Filter OAuth extractors based on hostname category
        if is_oauth_variable and not oauth_allowed:
            continue  # Skip OAuth extractor for non-auth hostnames
        
        # Check if we should skip this extractor (first_occurrence mode)
        if mode == "first_occurrence" and extracted_variables is not None:
            # For OAuth variables on auth hostnames, still allow multiple occurrences
            # For non-OAuth variables, only add first occurrence
            if not is_oauth_variable:
                if variable_name in extracted_variables:
                    continue  # Skip - already have an extractor for this variable
        
        extractor = _create_extractor_element(var_config)
        if extractor is not None:
            extractors.append(extractor)
            # Track that we've added an extractor for this variable
            if extracted_variables is not None:
                extracted_variables.add(variable_name)
    
    return extractors
