#!/usr/bin/env python3
"""
Import RF Device Signatures into Database

Imports device signatures from downloaded datasets:
1. Flipper Zero .sub file collections (with automatic labeling from file paths)
2. RTL_433 protocol definitions
3. Creates devices and signatures records for matching engine

Usage:
    python3 scripts/import_signatures_to_db.py --flipper --rtl433 --limit 10000
"""

import sys
import os
import json
import re
import hashlib
from pathlib import Path
from typing import Dict, List, Optional, Tuple
from dataclasses import dataclass

# Add project root to path
sys.path.insert(0, str(Path(__file__).parent.parent))

from sqlalchemy.orm import Session
from loguru import logger

from src.database.connection import get_session
from src.database.models import Device, Signature, FlipperSignature
from src.parser.sub_parser import parse_sub_file
from src.parser.metadata import SignalMetadata

# Compatibility alias
SessionLocal = get_session

# =============================================================================
# CONFIGURATION
# =============================================================================

DATA_DIR = Path(__file__).parent.parent / "data"

FLIPPER_DATASETS = [
    DATA_DIR / "rf_test_datasets" / "UberGuidoZ_Flipper",
    DATA_DIR / "test_known_devices",
]

RTL433_REPO = DATA_DIR / "test_rtl433_real"

# Device type mapping from file paths
DEVICE_TYPE_PATTERNS = {
    r"weather[_\s]?station": "Weather Sensor",
    r"garage[_\s]?door": "Garage Door Opener",
    r"gate[_\s]?opener": "Gate Opener",
    r"doorbell": "Doorbell",
    r"door[_\s]?sensor": "Door Sensor",
    r"motion[_\s]?sensor": "Motion Sensor",
    r"smoke[_\s]?detector": "Smoke Detector",
    r"remote[_\s]?control": "Remote Control",
    r"car[_\s]?key": "Car Key Fob",
    r"tpms": "TPMS",
    r"tire[_\s]?pressure": "TPMS",
    r"security": "Security System",
    r"alarm": "Alarm System",
}

MANUFACTURER_PATTERNS = {
    r"chamberlain": "Chamberlain",
    r"liftmaster": "LiftMaster",
    r"genie": "Genie",
    r"linear": "Linear",
    r"oregon[_\s]?scientific": "Oregon Scientific",
    r"acurite": "AcuRite",
    r"lacrosse": "LaCrosse",
    r"ambient[_\s]?weather": "Ambient Weather",
    r"honeywell": "Honeywell",
    r"nest": "Nest",
    r"ring": "Ring",
    r"yale": "Yale",
    r"toyota": "Toyota",
    r"schrader": "Schrader",
    r"nice": "NICE",
    r"came": "CAME",
    r"bft": "BFT",
}


# =============================================================================
# DATA CLASSES
# =============================================================================

@dataclass
class ExtractedDeviceInfo:
    """Device information extracted from file path and content"""
    device_type: Optional[str] = None
    manufacturer: Optional[str] = None
    model: Optional[str] = None
    confidence: float = 0.5  # Confidence in automatic extraction


# =============================================================================
# FILE PATH ANALYSIS
# =============================================================================

def extract_device_info_from_path(file_path: Path) -> ExtractedDeviceInfo:
    """
    Extract device information from file path structure

    Example paths:
        - Weather_stations/Oregon_Scientific/THGN123N.sub
        - Garage_Door_Openers/Chamberlain/model_891LM.sub
        - UberGuidoZ/Sub-GHz/Vehicles/Tesla/Model_3_Charge_Port.sub
    """
    info = ExtractedDeviceInfo()

    path_str = str(file_path).lower()
    parts = file_path.parts

    # Extract device type
    for pattern, device_type in DEVICE_TYPE_PATTERNS.items():
        if re.search(pattern, path_str, re.IGNORECASE):
            info.device_type = device_type
            info.confidence = 0.7
            break

    # Extract manufacturer
    for pattern, manufacturer in MANUFACTURER_PATTERNS.items():
        if re.search(pattern, path_str, re.IGNORECASE):
            info.manufacturer = manufacturer
            info.confidence = 0.8
            break

    # Extract model from filename (heuristic: uppercase letters + numbers)
    filename = file_path.stem  # Without extension
    model_match = re.search(r'([A-Z0-9]{3,}[-_]?[A-Z0-9]*)', filename)
    if model_match:
        info.model = model_match.group(1).replace('_', ' ')
        info.confidence = max(info.confidence, 0.6)

    return info


# =============================================================================
# FLIPPER ZERO SIGNATURE IMPORT
# =============================================================================

def import_flipper_signatures(
    db: Session,
    limit: Optional[int] = None,
    skip_existing: bool = True
) -> Tuple[int, int, int]:
    """
    Import Flipper Zero .sub files as device signatures

    Returns:
        (devices_created, signatures_created, files_skipped)
    """
    devices_created = 0
    signatures_created = 0
    files_skipped = 0

    # Collect all .sub files
    sub_files: List[Path] = []
    for dataset_dir in FLIPPER_DATASETS:
        if not dataset_dir.exists():
            logger.warning(f"Dataset not found: {dataset_dir}")
            continue

        logger.info(f"Scanning: {dataset_dir}")
        sub_files.extend(dataset_dir.rglob("*.sub"))

    logger.info(f"Found {len(sub_files)} .sub files across all datasets")

    if limit:
        sub_files = sub_files[:limit]
        logger.info(f"Limited to {limit} files")

    # Process each file
    for idx, sub_file in enumerate(sub_files, 1):
        if idx % 100 == 0:
            logger.info(f"Progress: {idx}/{len(sub_files)} ({idx/len(sub_files)*100:.1f}%)")

        try:
            # Compute file hash for deduplication
            file_hash = hashlib.sha256(sub_file.read_bytes()).hexdigest()

            # Check if already imported
            if skip_existing:
                existing = db.query(FlipperSignature).filter_by(
                    file_hash=file_hash
                ).first()
                if existing:
                    files_skipped += 1
                    continue

            # Parse .sub file
            try:
                metadata = parse_sub_file(str(sub_file))
            except Exception as e:
                logger.debug(f"Parse failed: {sub_file.name} - {e}")
                files_skipped += 1
                continue

            # Extract device info from path
            device_info = extract_device_info_from_path(sub_file)

            # Skip if no useful information extracted
            if not device_info.device_type and not metadata.protocol:
                files_skipped += 1
                continue

            # Create or find device
            device_query = db.query(Device)

            if device_info.manufacturer and device_info.model:
                device_query = device_query.filter_by(
                    manufacturer=device_info.manufacturer,
                    model=device_info.model
                )
            elif metadata.protocol:
                # Use protocol as fallback identifier
                device_query = device_query.filter_by(
                    protocol=metadata.protocol,
                    typical_frequency=metadata.frequency
                )
            else:
                files_skipped += 1
                continue

            device = device_query.first()

            if not device:
                # Create new device
                device = Device(
                    manufacturer=device_info.manufacturer or "Unknown",
                    model=device_info.model or metadata.protocol or "Unknown",
                    device_type=device_info.device_type or "Unknown",
                    typical_frequency=metadata.frequency,
                    protocol=metadata.protocol,
                    bit_length=metadata.bit_length,
                    encoding=metadata.file_format if metadata.file_format else None,
                    source="flipper",
                    verified=False  # Needs community verification
                )
                db.add(device)
                db.flush()  # Get device ID
                devices_created += 1

            # Extract timing statistics for signature matching
            timing_min = None
            timing_max = None

            if metadata.timing_element:
                # KEY format: use timing_element with tolerance
                timing_min = int(metadata.timing_element * 0.8)
                timing_max = int(metadata.timing_element * 1.2)
            elif metadata.raw_data and len(metadata.raw_data) > 0:
                # RAW format: calculate from pulse statistics
                pulse_widths = [abs(p) for p in metadata.raw_data]
                if pulse_widths:
                    timing_min = int(min(pulse_widths) * 0.9)  # 10% tolerance
                    timing_max = int(max(pulse_widths) * 1.1)

            # Generate bit_mask for KEY format (all bits significant by default)
            bit_mask = None
            if metadata.key_data:
                # Create mask with all bits set to 1 (all bits are significant)
                bit_mask = bytes([0xFF] * len(metadata.key_data))

            # Create signature with enhanced metadata
            signature = Signature(
                device_id=device.id,
                protocol=metadata.protocol,
                frequency=metadata.frequency,
                modulation=str(metadata.modulation).replace("Modulation.", "") if metadata.modulation else None,
                bit_pattern=metadata.key_data if metadata.key_data else None,
                bit_mask=bit_mask,
                timing_min=timing_min,
                timing_max=timing_max,
                weight=device_info.confidence,
                source="flipper"
            )
            db.add(signature)
            signatures_created += 1

            # Create Flipper-specific record
            flipper_sig = FlipperSignature(
                signature_id=None,  # Will be set after signature is committed
                file_hash=file_hash,
                file_path=str(sub_file.relative_to(DATA_DIR)),
                preset=metadata.preset if hasattr(metadata, 'preset') else None,
                raw_format=metadata.file_format
            )
            db.add(flipper_sig)

            # Commit in batches
            if idx % 50 == 0:
                try:
                    db.commit()
                except Exception as e:
                    logger.debug(f"Batch commit error: {e}")
                    db.rollback()

        except Exception as e:
            logger.error(f"Error processing {sub_file}: {e}")
            db.rollback()  # Rollback after error to continue with next file
            files_skipped += 1
            continue

    # Final commit
    try:
        db.commit()
    except Exception as e:
        logger.debug(f"Final commit error: {e}")
        db.rollback()

    logger.info(f"Import complete:")
    logger.info(f"  - Devices created: {devices_created}")
    logger.info(f"  - Signatures created: {signatures_created}")
    logger.info(f"  - Files skipped: {files_skipped}")

    return devices_created, signatures_created, files_skipped


# =============================================================================
# RTL_433 PROTOCOL IMPORT
# =============================================================================

def import_rtl433_protocols(db: Session) -> Tuple[int, int]:
    """
    Import RTL_433 protocol definitions from test files

    RTL_433 test files include .json metadata with:
    - Protocol name and number
    - Expected output fields
    - Device information

    Returns:
        (devices_created, signatures_created)
    """
    devices_created = 0
    signatures_created = 0

    if not RTL433_REPO.exists():
        logger.warning(f"RTL_433 test repo not found: {RTL433_REPO}")
        return 0, 0

    # Find all .json files (contain protocol metadata)
    json_files = list(RTL433_REPO.rglob("*.json"))
    logger.info(f"Found {len(json_files)} RTL_433 test .json files")

    protocol_map = {}  # protocol_name -> metadata

    for json_file in json_files:
        try:
            with open(json_file, 'r') as f:
                data = json.load(f)

            # RTL_433 .json files contain decoded device information
            if isinstance(data, dict) and 'model' in data:
                model = data.get('model')
                protocol_id = data.get('protocol', None)

                # Extract device info
                manufacturer = None
                if '-' in model:
                    parts = model.split('-', 1)
                    manufacturer = parts[0].strip()
                    model_name = parts[1].strip()
                else:
                    model_name = model

                # Frequency from file path or metadata
                freq_match = re.search(r'(\d+)M', str(json_file))
                frequency = int(freq_match.group(1)) * 1000000 if freq_match else 433920000

                # Store protocol info
                key = (manufacturer or "Unknown", model_name)
                if key not in protocol_map:
                    protocol_map[key] = {
                        'manufacturer': manufacturer or "Unknown",
                        'model': model_name,
                        'frequency': frequency,
                        'protocol': model,
                        'protocol_id': protocol_id,
                        'device_type': 'Weather Sensor',  # Most RTL_433 devices
                        'example_data': data
                    }

        except Exception as e:
            logger.debug(f"Error parsing {json_file}: {e}")
            continue

    logger.info(f"Extracted {len(protocol_map)} unique RTL_433 protocols")

    # Create devices and signatures
    for (manufacturer, model_name), proto_info in protocol_map.items():
        # Check if device exists
        device = db.query(Device).filter_by(
            manufacturer=manufacturer,
            model=model_name
        ).first()

        if not device:
            device = Device(
                manufacturer=manufacturer,
                model=model_name,
                device_type=proto_info['device_type'],
                typical_frequency=proto_info['frequency'],
                protocol=proto_info['protocol'],
                source="rtl433",
                verified=True  # RTL_433 protocols are trusted
            )
            db.add(device)
            db.flush()
            devices_created += 1

        # Create signature
        signature = Signature(
            device_id=device.id,
            protocol=proto_info['protocol'],
            frequency=proto_info['frequency'],
            weight=0.95  # High confidence for RTL_433
        )
        db.add(signature)
        signatures_created += 1

    db.commit()

    logger.info(f"RTL_433 import complete:")
    logger.info(f"  - Devices created: {devices_created}")
    logger.info(f"  - Signatures created: {signatures_created}")

    return devices_created, signatures_created


# =============================================================================
# MAIN
# =============================================================================

def main():
    import argparse

    parser = argparse.ArgumentParser(description="Import RF device signatures into database")
    parser.add_argument("--flipper", action="store_true", help="Import Flipper Zero signatures")
    parser.add_argument("--rtl433", action="store_true", help="Import RTL_433 protocols")
    parser.add_argument("--limit", type=int, default=None, help="Limit number of Flipper files to process")
    parser.add_argument("--skip-existing", action="store_true", default=True, help="Skip already imported files")
    parser.add_argument("--all", action="store_true", help="Import all sources")

    args = parser.parse_args()

    if not any([args.flipper, args.rtl433, args.all]):
        parser.print_help()
        print("\nError: Must specify at least one import source (--flipper, --rtl433, or --all)")
        sys.exit(1)

    # Setup logging
    logger.remove()
    logger.add(sys.stderr, level="INFO")
    logger.add("logs/import_signatures.log", rotation="10 MB", level="DEBUG")

    logger.info("=" * 80)
    logger.info("RF Device Signature Import")
    logger.info("=" * 80)

    # Connect to database
    db = SessionLocal()

    try:
        # Import Flipper signatures
        if args.flipper or args.all:
            logger.info("\n[1/2] Importing Flipper Zero signatures...")
            flipper_stats = import_flipper_signatures(
                db,
                limit=args.limit,
                skip_existing=args.skip_existing
            )
            logger.info(f"Flipper import: {flipper_stats[0]} devices, {flipper_stats[1]} signatures")

        # Import RTL_433 protocols
        if args.rtl433 or args.all:
            logger.info("\n[2/2] Importing RTL_433 protocols...")
            rtl_stats = import_rtl433_protocols(db)
            logger.info(f"RTL_433 import: {rtl_stats[0]} devices, {rtl_stats[1]} signatures")

        logger.info("\n" + "=" * 80)
        logger.info("Import complete!")
        logger.info("=" * 80)

        # Summary
        total_devices = db.query(Device).count()
        total_signatures = db.query(Signature).count()

        logger.info(f"\nDatabase totals:")
        logger.info(f"  - Total devices: {total_devices}")
        logger.info(f"  - Total signatures: {total_signatures}")

    finally:
        db.close()


if __name__ == "__main__":
    main()