#!/usr/bin/env python3
"""
Knowledge Graph Data Import Script for Tourism Recommender
Imports CSV data into Neo4j and validates relationships
"""

import pandas as pd
import numpy as np
from neo4j import GraphDatabase
import logging
from typing import Dict, List, Tuple
import os

# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

class TourismKGImporter:
    def __init__(self, uri=None, user=None, password=None):
        # Đọc từ env (NEO4J_URI/USER/PASSWORD) — cần thiết khi import lên
        # AuraDB (neo4j+s://...); mặc định rơi về Neo4j local trong Docker.
        uri = uri or os.getenv("NEO4J_URI", "bolt://localhost:7687")
        user = user or os.getenv("NEO4J_USER", "neo4j")
        password = password or os.getenv("NEO4J_PASSWORD", "tourism123")
        logger.info(f"Connecting to Neo4j at: {uri}")
        self.driver = GraphDatabase.driver(uri, auth=(user, password))
        
    def close(self):
        self.driver.close()
    
    def clear_database(self):
        """Clear all nodes and relationships"""
        with self.driver.session() as session:
            session.run("MATCH (n) DETACH DELETE n")
            logger.info("Database cleared")
    
    def create_constraints_and_indexes(self):
        """Create constraints and indexes for better performance"""
        with self.driver.session() as session:
            constraints = [
                "CREATE CONSTRAINT location_id IF NOT EXISTS FOR (l:Location) REQUIRE l.id IS UNIQUE",
                "CREATE CONSTRAINT ward_id IF NOT EXISTS FOR (w:Ward) REQUIRE w.id IS UNIQUE", 
                "CREATE CONSTRAINT province_id IF NOT EXISTS FOR (p:Province) REQUIRE p.id IS UNIQUE",
                "CREATE CONSTRAINT sub_unit_id IF NOT EXISTS FOR (s:SubUnit) REQUIRE s.id IS UNIQUE",
                "CREATE INDEX location_name IF NOT EXISTS FOR (l:Location) ON (l.name)",
                "CREATE INDEX ward_name IF NOT EXISTS FOR (w:Ward) ON (w.name)",
                "CREATE INDEX sub_unit_name IF NOT EXISTS FOR (s:SubUnit) ON (s.name)"
            ]
            
            for constraint in constraints:
                try:
                    session.run(constraint)
                    logger.info(f"Created: {constraint}")
                except Exception as e:
                    logger.warning(f"Constraint/Index already exists or failed: {e}")
    
    def load_provinces(self, provinces_file="Provinces.csv"):
        """Load provinces data"""
        logger.info("Loading provinces...")
        df = pd.read_csv(provinces_file)
        
        with self.driver.session() as session:
            for _, row in df.iterrows():
                session.run("""
                    MERGE (p:Province {id: $province_id})
                    SET p.old_name = $old_name,
                        p.new_name = $new_name,
                        p.name = $new_name
                """, 
                province_id=int(row['province_id']),
                old_name=row['old_province_name'],
                new_name=row['new_province_name']
                )
        
        logger.info(f"Loaded {len(df)} provinces")
    
    def load_wards(self, wards_file="Wards.csv"):
        """Load wards data and create relationships with provinces"""
        logger.info("Loading wards...")
        df = pd.read_csv(wards_file)
        
        with self.driver.session() as session:
            for _, row in df.iterrows():
                # Create ward node
                session.run("""
                    MERGE (w:Ward {id: $ward_id})
                    SET w.old_type = $old_type,
                        w.old_name = $old_name,
                        w.new_type = $new_type,
                        w.new_name = $new_name,
                        w.name = $new_name,
                        w.district_city = $district_city
                """,
                ward_id=int(row['ward_id']),
                old_type=row['old_type'],
                old_name=row['old_name'],
                new_type=row['new_type'],
                new_name=row['new_name'],
                district_city=row['district_city']
                )
                
                # Create relationship with province
                session.run("""
                    MATCH (w:Ward {id: $ward_id})
                    MATCH (p:Province {id: $province_id})
                    MERGE (w)-[:BELONGS_TO]->(p)
                """,
                ward_id=int(row['ward_id']),
                province_id=int(row['province_id'])
                )
        
        logger.info(f"Loaded {len(df)} wards")
    
    def load_tourism_locations(self, tourism_file="Tourism_FInal.csv"):
        """Load tourism locations and create relationships with wards"""
        logger.info("Loading tourism locations...")
        df = pd.read_csv(tourism_file)
        
        with self.driver.session() as session:
            for _, row in df.iterrows():
                # Create location node
                session.run("""
                    MERGE (l:Location {id: $location_id})
                    SET l.name = $name,
                        l.lat = $lat,
                        l.lng = $lng,
                        l.coordinates = point({latitude: $lat, longitude: $lng})
                """,
                location_id=f"loc_{row['ward_id']}_{hash(row['name']) % 10000}",
                name=row['name'],
                lat=float(row['lat']),
                lng=float(row['lng'])
                )
                
                # Create relationship with ward
                session.run("""
                    MATCH (l:Location {id: $location_id})
                    MATCH (w:Ward {id: $ward_id})
                    MERGE (l)-[:LOCATED_IN]->(w)
                """,
                location_id=f"loc_{row['ward_id']}_{hash(row['name']) % 10000}",
                ward_id=int(row['ward_id'])
                )
        
        logger.info(f"Loaded {len(df)} tourism locations")
    
    def load_tourism_locations_xlsx(self, tourism_file="datasets/dulich.xlsx"):
        """Load tourism POIs from dulich.xlsx (the file actually shipped in the
        repo) and link them to wards.

        Column mapping (Vietnamese headers in dulich.xlsx):
            'Tên địa điểm'        -> name
            'Mô tả ngắn'          -> description
            'Địa chỉ'             -> address
            'Số điện thoại'       -> phone
            'Giờ mở cửa'          -> opening_hours
            'Id_ xã phường cũ'    -> ward_id (links to Ward.id)
            'Lat' / 'Lng'         -> coordinates

        Uses a stable, collision-resistant id (sha1 of ward_id + name) instead of
        the original ``hash(name) % 10000`` which could silently overwrite POIs.
        """
        import hashlib

        logger.info("Loading tourism locations from xlsx...")
        df = pd.read_excel(tourism_file)

        def col(row, *names, default=None):
            for n in names:
                if n in row and pd.notna(row[n]):
                    return row[n]
            return default

        import re as _re

        def fix_coord(value, limit):
            """Best-effort coordinate parser for messy xlsx cells.

            Handles: plain numbers; numbers that lost the decimal separator
            (106993595 -> 106.993595); free text like '~11.1848° N' or
            '10°57′53.28″N ≈ 10.9648' (takes the LAST decimal number found);
            trailing dots ('106.7930.'). Returns None when no usable number.
            """
            if value is None:
                return None
            s = str(value)
            nums = _re.findall(r"[-+]?\d+(?:[.,]\d+)?", s)
            if not nums:
                return None
            try:
                f = float(nums[-1].replace(",", "."))
            except ValueError:
                return None
            guard = 0
            while abs(f) > limit and guard < 12:
                f /= 10.0
                guard += 1
            return f if abs(f) <= limit else None

        loaded_xy, loaded_noxy, skipped = 0, 0, 0
        with self.driver.session() as session:
            for _, row in df.iterrows():
                name = col(row, "Tên địa điểm", "name")
                ward_raw = col(row, "Id_ xã phường cũ", "ward_id")
                lat = col(row, "Lat", "lat")
                lng = col(row, "Lng", "lng")
                if not name or ward_raw is None:
                    skipped += 1
                    continue
                try:
                    ward_id = int(ward_raw)
                except (ValueError, TypeError):
                    skipped += 1
                    continue

                # Coordinates are OPTIONAL: most dulich.xlsx rows lack them or
                # carry free text. POIs without coordinates are still imported
                # (they just won't take part in NEAR relationships).
                lat_f = fix_coord(lat, 90.0)
                lng_f = fix_coord(lng, 180.0)
                has_xy = (lat_f is not None and lng_f is not None
                          and 0 < lat_f < 30 and 100 < lng_f < 115)  # Vietnam bounds

                loc_hash = hashlib.sha1(f"{ward_id}|{name}".encode("utf-8")).hexdigest()[:12]
                location_id = f"loc_{ward_id}_{loc_hash}"
                props = dict(
                    location_id=location_id,
                    name=str(name),
                    description=str(col(row, "Mô tả ngắn", default="")),
                    address=str(col(row, "Địa chỉ", default="")),
                    phone=str(col(row, "Số điện thoại", default="")),
                    opening_hours=str(col(row, "Giờ mở cửa", default="")),
                )

                if has_xy:
                    session.run(
                        """
                        MERGE (l:Location {id: $location_id})
                        SET l.name = $name,
                            l.description = $description,
                            l.address = $address,
                            l.phone = $phone,
                            l.opening_hours = $opening_hours,
                            l.lat = $lat,
                            l.lng = $lng,
                            l.coordinates = point({latitude: $lat, longitude: $lng})
                        """,
                        lat=lat_f, lng=lng_f, **props,
                    )
                    loaded_xy += 1
                else:
                    session.run(
                        """
                        MERGE (l:Location {id: $location_id})
                        SET l.name = $name,
                            l.description = $description,
                            l.address = $address,
                            l.phone = $phone,
                            l.opening_hours = $opening_hours
                        """,
                        **props,
                    )
                    loaded_noxy += 1

                session.run(
                    """
                    MATCH (l:Location {id: $location_id})
                    MATCH (w:Ward {id: $ward_id})
                    MERGE (l)-[:LOCATED_IN]->(w)
                    """,
                    location_id=location_id,
                    ward_id=ward_id,
                )

        logger.info(
            f"Loaded {loaded_xy + loaded_noxy} tourism locations from xlsx "
            f"({loaded_xy} with coordinates, {loaded_noxy} without, {skipped} skipped)"
        )

    def load_sub_units(self, sub_units_file="mapped_sub_units.csv"):
        """Load sub-units data and create relationships with wards"""
        logger.info("Loading sub-units...")
        df = pd.read_csv(sub_units_file)
        
        with self.driver.session() as session:
            skipped_rows = 0
            for _, row in df.iterrows():
                try:
                    ward_id = int(row['ward_id'])
                    
                    # Create sub-unit node
                    session.run("""
                        MERGE (s:SubUnit {id: $sub_unit_id})
                        SET s.old_type = $old_type,
                            s.old_name = $old_name,
                            s.new_type = $new_type,
                            s.new_name = $new_name,
                            s.name = $new_name
                    """,
                    sub_unit_id=f"sub_{ward_id}_{hash(row['new_name']) % 10000}",
                    old_type=row['old_type'],
                    old_name=row['old_name'],
                    new_type=row['new_type'],
                    new_name=row['new_name']
                    )
                    
                    # Create relationship with ward
                    session.run("""
                        MATCH (s:SubUnit {id: $sub_unit_id})
                        MATCH (w:Ward {id: $ward_id})
                        MERGE (s)-[:PART_OF]->(w)
                    """,
                    sub_unit_id=f"sub_{ward_id}_{hash(row['new_name']) % 10000}",
                    ward_id=ward_id
                    )
                except (ValueError, TypeError) as e:
                    logger.warning(f"Skipping row with invalid ward_id: {row['ward_id']} - {e}")
                    skipped_rows += 1
                    continue
        
        logger.info(f"Loaded {len(df) - skipped_rows} sub-units (skipped {skipped_rows} invalid rows)")
    
    def create_proximity_relationships(self, distance_km=5.0):
        """Create NEAR relationships between locations within specified distance"""
        logger.info(f"Creating proximity relationships (within {distance_km}km)...")
        
        with self.driver.session() as session:
            # Use Neo4j's spatial functions to find nearby locations
            session.run("""
                MATCH (l1:Location), (l2:Location)
                WHERE l1 <> l2 
                AND point.distance(l1.coordinates, l2.coordinates) < $distance_meters
                MERGE (l1)-[:NEAR {distance: point.distance(l1.coordinates, l2.coordinates)}]->(l2)
            """, distance_meters=distance_km * 1000)
        
        logger.info("Proximity relationships created")
    
    def validate_and_clean_orphaned_nodes(self):
        """Find and optionally remove nodes without relationships"""
        logger.info("Validating node relationships...")
        
        with self.driver.session() as session:
            # Find orphaned locations (no relationships)
            result = session.run("""
                MATCH (l:Location)
                WHERE NOT (l)-[]-()
                RETURN l.id as location_id, l.name as name
            """)
            
            orphaned_locations = [record for record in result]
            if orphaned_locations:
                logger.warning(f"Found {len(orphaned_locations)} orphaned locations:")
                for loc in orphaned_locations[:10]:  # Show first 10
                    logger.warning(f"  - {loc['name']} (ID: {loc['location_id']})")
            
            # Find orphaned wards
            result = session.run("""
                MATCH (w:Ward)
                WHERE NOT (w)-[]-()
                RETURN w.id as ward_id, w.name as name
            """)
            
            orphaned_wards = [record for record in result]
            if orphaned_wards:
                logger.warning(f"Found {len(orphaned_wards)} orphaned wards:")
                for ward in orphaned_wards[:10]:
                    logger.warning(f"  - {ward['name']} (ID: {ward['ward_id']})")
            
            return len(orphaned_locations), len(orphaned_wards)
    
    def get_database_stats(self):
        """Get database statistics"""
        with self.driver.session() as session:
            stats = {}
            
            # Count nodes by type
            result = session.run("MATCH (n:Location) RETURN count(n) as count")
            stats['locations'] = result.single()['count']
            
            result = session.run("MATCH (n:Ward) RETURN count(n) as count")
            stats['wards'] = result.single()['count']
            
            result = session.run("MATCH (n:Province) RETURN count(n) as count")
            stats['provinces'] = result.single()['count']
            
            result = session.run("MATCH (n:SubUnit) RETURN count(n) as count")
            stats['sub_units'] = result.single()['count']
            
            # Count relationships by type
            result = session.run("MATCH ()-[r:LOCATED_IN]->() RETURN count(r) as count")
            stats['located_in_rels'] = result.single()['count']
            
            result = session.run("MATCH ()-[r:BELONGS_TO]->() RETURN count(r) as count")
            stats['belongs_to_rels'] = result.single()['count']
            
            result = session.run("MATCH ()-[r:NEAR]->() RETURN count(r) as count")
            stats['near_rels'] = result.single()['count']
            
            result = session.run("MATCH ()-[r:PART_OF]->() RETURN count(r) as count")
            stats['part_of_rels'] = result.single()['count']
            
            return stats
    
    def run_full_import(self, clear_existing=True):
        """Run complete data import process"""
        logger.info("Starting full knowledge graph import...")
        
        if clear_existing:
            self.clear_database()
        
        # Create constraints and indexes
        self.create_constraints_and_indexes()
        
        # Import data in order
        self.load_provinces("datasets/Provinces.csv")
        self.load_wards("datasets/Wards.csv")
        # Prefer the legacy Tourism_FInal.csv if present, otherwise fall back to
        # the dulich.xlsx file actually shipped in the repo.
        if os.path.exists("datasets/Tourism_FInal.csv"):
            self.load_tourism_locations("datasets/Tourism_FInal.csv")
        else:
            logger.info("Tourism_FInal.csv not found; loading POIs from datasets/dulich.xlsx")
            self.load_tourism_locations_xlsx("datasets/dulich.xlsx")
        self.load_sub_units("datasets/mapped_sub_units.csv")
        
        # Create proximity relationships
        self.create_proximity_relationships()
        
        # Validate data
        orphaned_locs, orphaned_wards = self.validate_and_clean_orphaned_nodes()
        
        # Get final stats
        stats = self.get_database_stats()
        
        logger.info("Import completed!")
        logger.info(f"Database statistics:")
        for key, value in stats.items():
            logger.info(f"  {key}: {value}")
        
        if orphaned_locs > 0 or orphaned_wards > 0:
            logger.warning(f"Found {orphaned_locs} orphaned locations and {orphaned_wards} orphaned wards")
            logger.warning("Consider reviewing data quality or relationship logic")
        
        return stats

def main():
    """Main execution function"""
    importer = TourismKGImporter()
    
    try:
        # Test connection
        with importer.driver.session() as session:
            result = session.run("RETURN 'Connection successful' as message")
            logger.info(result.single()['message'])
        
        # Run import
        stats = importer.run_full_import(clear_existing=True)
        
        logger.info("Knowledge graph import completed successfully!")
        
    except Exception as e:
        logger.error(f"Import failed: {e}")
        raise
    finally:
        importer.close()

if __name__ == "__main__":
    main()