# internal_geocoding.py
# ------------------------------------------------------------
# Internal geocoding using PostgreSQL pg_trgm fuzzy matching
# Searches AdrPoint and AdrStreet models using trigram similarity
# Requires: CREATE EXTENSION pg_trgm; 
# Recommended: GiST indexes with gist_trgm_ops for optimal ORDER BY <-> performance
# ------------------------------------------------------------
import logging
import asyncio
from concurrent.futures import ThreadPoolExecutor
from typing import Optional, Tuple, List, Dict
from django.db import connection

from common.models import AdrPoint, AdrStreet

logger = logging.getLogger(__name__)

# Thread pool for parallel DB queries - limit to avoid connection pool exhaustion
_db_executor = ThreadPoolExecutor(max_workers=5)

# Default score threshold (0.0 to 100.0) - addresses with score below this will fallback to Photon
# pg_trgm similarity scores range from 0.0 to 1.0 (scaled to 0-100 in results)
DEFAULT_SCORE_THRESHOLD = 30.0


def normalize_address(address: str) -> str:
    """Normalize an address string for better matching"""
    if not address:
        return ""
    # Convert to lowercase and strip whitespace
    address = address.lower().strip()
    # Remove extra spaces
    address = " ".join(address.split())
    return address


def search_address_points(query: str, limit: int = 5) -> List[Tuple[AdrPoint, float]]:
    """
    Search for address points in the database using PostgreSQL trigram distance.
    Uses the <-> operator which is optimized for ORDER BY LIMIT queries.
    
    Runs two separate queries (FR and NL) and merges results for optimal index usage.
    
    Args:
        query: Address query string
        limit: Maximum number of results to return
    
    Returns:
        List of tuples (AdrPoint, score) ordered by score descending
    """
    normalized_query = normalize_address(query)
    
    if not normalized_query:
        return []
    
    # Run two separate queries for FR and NL to properly use GiST indexes
    # Using LEAST() on two columns prevents index usage
    sql_fr = """
        SELECT id, 1 - (full_address_fr <-> %s) as score
        FROM common_adrpoint
        WHERE lon IS NOT NULL AND lat IS NOT NULL AND full_address_fr IS NOT NULL
        ORDER BY full_address_fr <-> %s
        LIMIT %s;
    """
    
    sql_nl = """
        SELECT id, 1 - (full_address_nl <-> %s) as score
        FROM common_adrpoint
        WHERE lon IS NOT NULL AND lat IS NOT NULL AND full_address_nl IS NOT NULL
        ORDER BY full_address_nl <-> %s
        LIMIT %s;
    """
    
    with connection.cursor() as cursor:
        # Query FR
        cursor.execute(sql_fr, [normalized_query, normalized_query, limit])
        rows_fr = cursor.fetchall()
        
        # Query NL
        cursor.execute(sql_nl, [normalized_query, normalized_query, limit])
        rows_nl = cursor.fetchall()
    
    # Merge results, keeping best score per id (scaling pg_trgm scores to 0-100)
    scores_by_id: Dict[int, float] = {}
    for row in rows_fr:
        scores_by_id[row[0]] = row[1] * 100.0
    for row in rows_nl:
        val = row[1] * 100.0
        if row[0] in scores_by_id:
            scores_by_id[row[0]] = max(scores_by_id[row[0]], val)
        else:
            scores_by_id[row[0]] = val
    
    if not scores_by_id:
        return []
    
    # Get top N by score
    top_ids = sorted(scores_by_id.keys(), key=lambda x: scores_by_id[x], reverse=True)[:limit]
    
    # Fetch the actual AdrPoint objects
    points = list(AdrPoint.objects.select_related(
        'street_name', 'municipality'
    ).filter(id__in=top_ids))
    
    # Build results maintaining the score order
    results = [(point, scores_by_id[point.id]) for point in points]
    results.sort(key=lambda x: x[1], reverse=True)
    
    return results


def search_streets(query: str, limit: int = 5) -> List[Tuple[AdrStreet, float]]:
    """
    Search for streets in the database using PostgreSQL trigram distance.
    Uses the <-> operator which is optimized for ORDER BY LIMIT queries.
    
    Runs two separate queries (FR and NL) and merges results for optimal index usage.
    
    Args:
        query: Street query string
        limit: Maximum number of results to return
    
    Returns:
        List of tuples (AdrStreet, score) ordered by score descending
    """
    normalized_query = normalize_address(query)
    
    if not normalized_query:
        return []
    
    # Run two separate queries for FR and NL to properly use GiST indexes
    sql_fr = """
        SELECT id, 1 - (full_street_fr <-> %s) as score
        FROM common_adrstreet
        WHERE lon IS NOT NULL AND lat IS NOT NULL AND full_street_fr IS NOT NULL
        ORDER BY full_street_fr <-> %s
        LIMIT %s;
    """
    
    sql_nl = """
        SELECT id, 1 - (full_street_nl <-> %s) as score
        FROM common_adrstreet
        WHERE lon IS NOT NULL AND lat IS NOT NULL AND full_street_nl IS NOT NULL
        ORDER BY full_street_nl <-> %s
        LIMIT %s;
    """
    
    with connection.cursor() as cursor:
        # Query FR
        cursor.execute(sql_fr, [normalized_query, normalized_query, limit])
        rows_fr = cursor.fetchall()
        
        # Query NL
        cursor.execute(sql_nl, [normalized_query, normalized_query, limit])
        rows_nl = cursor.fetchall()
    
    # Merge results, keeping best score per id (scaling pg_trgm scores to 0-100)
    scores_by_id: Dict[int, float] = {}
    for row in rows_fr:
        scores_by_id[row[0]] = row[1] * 100.0
    for row in rows_nl:
        val = row[1] * 100.0
        if row[0] in scores_by_id:
            scores_by_id[row[0]] = max(scores_by_id[row[0]], val)
        else:
            scores_by_id[row[0]] = val
    
    if not scores_by_id:
        return []
    
    # Get top N by score
    top_ids = sorted(scores_by_id.keys(), key=lambda x: scores_by_id[x], reverse=True)[:limit]
    
    # Fetch the actual AdrStreet objects
    streets = list(AdrStreet.objects.select_related('municipality').filter(id__in=top_ids))
    
    # Build results maintaining the score order
    results = [(street, scores_by_id[street.id]) for street in streets]
    results.sort(key=lambda x: x[1], reverse=True)
    
    return results


def search_streets_in_memory(query: str, limit: int = 5) -> List[Tuple[AdrStreet, float]]:
    """
    Alternative: Search streets using in-memory fuzzy matching with rapidfuzz.
    Useful when pg_trgm is not available or for small datasets (~3000 streets).
    
    Args:
        query: Street query string
        limit: Maximum number of results to return
    
    Returns:
        List of tuples (AdrStreet, score) ordered by score descending
    """
    try:
        from rapidfuzz import fuzz
    except ImportError:
        logger.warning("rapidfuzz not installed, falling back to pg_trgm")
        return search_streets(query, limit)
    
    normalized_query = normalize_address(query)
    
    if not normalized_query:
        return []
    
    # Load all streets in memory (only ~3000 records)
    streets = AdrStreet.objects.select_related('municipality').filter(
        lon__isnull=False,
        lat__isnull=False,
        full_street_fr__isnull=False
    )
    
    results = []
    for street in streets:
        # Calculate scores using rapidfuzz (already on 0-100 scale)
        score_fr = fuzz.token_set_ratio(normalized_query, street.full_street_fr.lower() if street.full_street_fr else "")
        score_nl = fuzz.token_set_ratio(normalized_query, street.full_street_nl.lower() if street.full_street_nl else "")
        best_score = max(score_fr, score_nl)
        
        if best_score > 10.0:
            results.append((street, best_score))
    
    # Sort by score descending and return top results
    results.sort(key=lambda x: x[1], reverse=True)
    return results[:limit]


def geocode_internal(
    address: str,
    threshold: float = DEFAULT_SCORE_THRESHOLD
) -> Tuple[Optional[float], Optional[float], str, float]:
    """
    Try to geocode an address using the internal database with PostgreSQL trigram matching.
    
    Args:
        address: Address string to geocode
        threshold: Minimum score threshold (0.0-1.0) to accept a match
    
    Returns:
        Tuple of (longitude, latitude, found_address, score)
        - On success: (lon, lat, formatted_address, score)
        - On no match: (None, None, "No match in internal DB", 0)
    """
    if not address or not address.strip():
        return (None, None, "Empty address", 0)
    
    # First, try to find address points (more precise)
    point_results = search_address_points(address, limit=1)
    
    if point_results:
        point, score = point_results[0]
        
        if score >= threshold:
            # Use the pre-computed full address or build it
            found_address = point.full_address_fr or ""
            if not found_address:
                street_name = point.street_name.name_fr if point.street_name else ""
                municipality_name = point.municipality.name if point.municipality else ""
                found_address = f"{point.number} {street_name}"
                if municipality_name:
                    found_address = f"{found_address}, {municipality_name}"
            
            logger.debug(f"Found address point: '{address}' -> '{found_address}' (score: {score:.2f})")
            return (point.lon, point.lat, found_address, score)
    
    # If no address point found, try streets
    street_results = search_streets(address, limit=1)
    
    if street_results:
        street, score = street_results[0]
        
        if score >= threshold:
            # Use the pre-computed full street name or build it
            found_address = street.full_street_fr or ""
            if not found_address:
                street_name = street.name_fr
                municipality_name = street.municipality.name if street.municipality else ""
                found_address = street_name
                if municipality_name:
                    found_address = f"{found_address}, {municipality_name}"
            
            logger.debug(f"Found street: '{address}' -> '{found_address}' (score: {score:.2f})")
            return (street.lon, street.lat, found_address, score)
    
    # No match found above threshold
    all_scores = [s for _, s in point_results + street_results]
    best_score_found = max(all_scores, default=0)
    logger.debug(f"No internal match for '{address}' (best score: {best_score_found:.2f})")
    return (None, None, "No match in internal DB", 0)


async def geocode_internal_batch(
    addresses: List[Tuple[int, str]],
    threshold: float = DEFAULT_SCORE_THRESHOLD,
    max_concurrent: int = 5,
) -> Dict[int, Tuple[Optional[float], Optional[float], str, float]]:
    """
    Geocode multiple addresses in parallel using the internal database.
    
    Args:
        addresses: List of (index, address) tuples
        threshold: Minimum score threshold (0.0-1.0) to accept a match
        max_concurrent: Maximum concurrent DB queries (default 5)
    
    Returns:
        Dict mapping index to (longitude, latitude, found_address, score)
    """
    if not addresses:
        return {}
    
    loop = asyncio.get_event_loop()
    
    # Create tasks for parallel execution
    async def geocode_one(idx: int, addr: str) -> Tuple[int, Tuple[Optional[float], Optional[float], str, float]]:
        """Geocode a single address and return (index, result)"""
        try:
            result = await loop.run_in_executor(
                _db_executor, 
                geocode_internal, 
                addr, 
                threshold
            )
            return (idx, result)
        except Exception as e:
            logger.warning(f"Internal geocoding error for index {idx}: {e}")
            return (idx, (None, None, f"Error: {e}", 0))
    
    # Process in chunks to limit concurrent DB connections
    results: Dict[int, Tuple[Optional[float], Optional[float], str, float]] = {}
    
    for i in range(0, len(addresses), max_concurrent):
        chunk = addresses[i:i + max_concurrent]
        tasks = [geocode_one(idx, addr) for idx, addr in chunk]
        chunk_results = await asyncio.gather(*tasks, return_exceptions=True)
        
        for res in chunk_results:
            if isinstance(res, Exception):
                logger.warning(f"Unexpected error in batch geocoding: {res}")
            else:
                idx, geocode_result = res
                results[idx] = geocode_result
    
    return results