import os
import pg8000.dbapi as pg
import pandas as pd
import uuid
import datetime
import sys

# Output log
log_file = open('rebuild_log.txt', 'w', encoding='utf-8')
sys.stdout = log_file

def log(msg):
    print(msg)
    log_file.flush()

def load_env():
    env_vars = {}
    try:
        with open('.env', 'r', encoding='utf-8') as f:
            for line in f:
                line = line.strip()
                if not line or line.startswith('#'):
                    continue
                if '=' in line:
                    k, v = line.split('=', 1)
                    env_vars[k.strip()] = v.strip()
    except Exception:
        pass
    return env_vars

env = load_env()
HOST = env.get("LIC_PG_HOST", "172.24.240.1")
PORT = int(env.get("LIC_PG_PORT", "5432"))
USER = env.get("LIC_PG_USER", "postgres")
PASSWORD = env.get("LIC_PG_PASSWORD", "")
DATABASE = env.get("LIC_PG_DATABASE", "licensing_risks")

CONN = None

def get_conn():
    global CONN
    if CONN is None:
        CONN = pg.connect(host=HOST, port=PORT, user=USER, password=PASSWORD, database=DATABASE)
        CONN.autocommit = True
    return CONN

def run():
    try:
        conn = get_conn()
        cursor = conn.cursor()
        
        # 1. Clear existing data
        log("Clearing existing data...")
        # Order matters for foreign keys if they exist
        cursor.execute("DELETE FROM region_theme_permits")
        cursor.execute("DELETE FROM permit_theme_rules")
        cursor.execute("DELETE FROM themes") 
        # cursor.execute("DELETE FROM region_themes") # If it exists? Check if it errors?
        try:
             cursor.execute("DELETE FROM region_themes")
        except:
             pass 
             
        log("Data cleared.")

        # 2. Get Region ID for "市级"
        log("Fetching Region ID for '市级'...")
        cursor.execute("SELECT id FROM regions WHERE name = '市级'")
        res = cursor.fetchone()
        if not res:
            log("Error: '市级' region not found!")
            return
        region_id = res[0]
        log(f"Region ID: {region_id}")

        # 3. Read Excel
        log("Reading Excel...")
        df = pd.read_excel('主题-事项绑定.xlsx', sheet_name='新主题事项表-12个')
        # Rename columns to standard
        df.rename(columns={'主题名称': 'theme_name', '风险提示事项名称': 'permit_name'}, inplace=True)
        # Fallback if old name is used
        if 'permit_name' not in df.columns and '事项名称' in df.columns:
            df.rename(columns={'事项名称': 'permit_name'}, inplace=True)
        # Forward fill theme_name (handle merged cells which come as NaN)
        df['theme_name'] = df['theme_name'].fillna(method='ffill')
        
        # Clean data
        df['theme_name'] = df['theme_name'].astype(str).str.strip()
        df['permit_name'] = df['permit_name'].astype(str).str.strip()
        
        # Drop duplicates to prevent unique constraint violations
        original_count = len(df)
        df.drop_duplicates(subset=['theme_name', 'permit_name'], inplace=True)
        log(f"Dropped {original_count - len(df)} duplicate rows.")
        
        log(f"Found {len(df)} unique rows.")

        # 4. Prepare Lookups
        log("Loading existing permits...")
        cursor.execute("SELECT name, id FROM permits")
        permit_map = {row[0].strip(): row[1] for row in cursor.fetchall()}
        log(f"Loaded {len(permit_map)} existing permits.")

        # 5. Process
        theme_map = {} # name -> uuid
        results = []

        for idx, row in df.iterrows():
            t_name = row['theme_name']
            p_name = row['permit_name']
            
            if pd.isna(t_name) or t_name == 'nan' or not t_name:
                results.append({'Theme': t_name, 'Permit': p_name, 'Status': 'Skipped (No Theme)'})
                continue
                
            if pd.isna(p_name) or p_name == 'nan' or not p_name:
                results.append({'Theme': t_name, 'Permit': p_name, 'Status': 'Skipped (No Permit Name)'})
                continue

            # Create Theme if needed
            if t_name not in theme_map:
                t_id = uuid.uuid4()
                cursor.execute("INSERT INTO themes (id, name) VALUES (%s, %s)", (t_id, t_name))
                theme_map[t_name] = t_id
            t_id = theme_map[t_name]

            # Bind Rule (name based)
            rule_id = uuid.uuid4()
            cursor.execute("""
                INSERT INTO permit_theme_rules (id, theme_id, permit_name, region_id, created_at)
                VALUES (%s, %s, %s, %s, NOW())
            """, (rule_id, t_id, p_name, region_id))

            # Bind Physical (id based)
            status = "Bound (Rule Only)"
            if p_name in permit_map:
                p_id = permit_map[p_name]
                # Ensure theme is registered in the region
                cursor.execute("""
                    INSERT INTO region_themes (region_id, theme_id)
                    VALUES (%s, %s)
                    ON CONFLICT DO NOTHING
                """, (region_id, t_id))
                
                cursor.execute("""
                    INSERT INTO region_theme_permits (region_id, theme_id, permit_id)
                    VALUES (%s, %s, %s)
                """, (region_id, t_id, p_id))
                status = "Bound (Success)"
            else:
                status = "Bound Rule, but Permit Not Found in DB"
            
            results.append({'Theme': t_name, 'Permit': p_name, 'Status': status})

        # 6. Save Report
        log("Saving report...")
        report_df = pd.DataFrame(results)
        try:
            report_df.to_excel('主题-事项绑定结果.xlsx', index=False)
        except PermissionError:
            log("File locked, saving to '主题-事项绑定结果_new.xlsx' instead.")
            report_df.to_excel('主题-事项绑定结果_new.xlsx', index=False)
        log("Done.")
        conn.close()

    except Exception as e:
        log(f"CRITICAL ERROR: {e}")
        import traceback
        traceback.print_exc(file=log_file)

if __name__ == '__main__':
    run()
