import os
import pg8000.dbapi as pg
from lawrisk.utils.env_loader import load_env
from lawrisk.services.licensing_repo import _parse_import_workbook, _ensure_risk

def patch_missing_risks():
    load_env()
    file_path = r"c:\Users\WIN10\Desktop\work\11th-week\法律风险提示-new\市监局-lawRisk-backend\市级初版-20251219\许可风险提示\101（生产科） 风险提示表(食品小作坊登记，市场监管部门) (1)_转自XLS.xlsx"
    target_sheet_name = "市级" # Will match partially
    target_permit_name = "食品小作坊登记证"
    
    with open(file_path, "rb") as f:
        content = f.read()
        
    conn_params = {
        "host": os.getenv("LIC_PG_HOST", "172.24.240.1"),
        "port": int(os.getenv("LIC_PG_PORT", "5432")),
        "user": os.getenv("LIC_PG_USER", "postgres"),
        "password": os.getenv("LIC_PG_PASSWORD", ""),
        "database": "licensing_risks",
    }
    
    try:
        conn = pg.connect(**conn_params)
        cur = conn.cursor()
        
        # 1. Get IDs
        cur.execute("SELECT id FROM regions WHERE name = '市级'")
        region_id = cur.fetchone()[0]
        cur.execute("SELECT id FROM permits WHERE name = '食品小作坊登记证'")
        permit_id = cur.fetchone()[0]
        print(f"Targeting: Region={region_id}, Permit={permit_id}")
        
        # 2. Parse Excel
        print("Parsing workbook...")
        parsed = _parse_import_workbook(content, "101.xlsx")
        sheets = parsed.get("sheets", {})
        
        target_rows = []
        for s_name, s_data in sheets.items():
            if target_sheet_name in s_name:
                print(f"Found match sheet: {s_name}")
                target_rows = s_data.get("rows", [])
                break
                
        if not target_rows:
            print("Target sheet not found!")
            return

        print(f"Processing {len(target_rows)} rows...")
        restored_count = 0
        
        for r in target_rows:
            # Check if it's the right permit
            if target_permit_name not in r.get("permit_name", ""):
                continue

            # Ensure risk exists in 'risks' table (idempotent)
            risk_id = _ensure_risk(
                conn,
                risk_content=r["risk_content"],
                legal_basis=r.get("legal_basis"),
                document_no=r.get("document_no"),
                summary=r.get("summary"),
                remark=r.get("remark")
            )
            
            serial = r.get("serial_number")
            
            # Upsert into region_permit_risks
            # This is the key patching step
            cur.execute("""
                INSERT INTO region_permit_risks (region_id, permit_id, risk_id, serial_number)
                VALUES (%s, %s, %s, %s)
                ON CONFLICT (region_id, permit_id, risk_id) 
                DO UPDATE SET serial_number = EXCLUDED.serial_number
            """, (region_id, permit_id, risk_id, serial))
            
            # Simple heuristic to say we 'restored' or updated it
            restored_count += 1
            if serial in [3, 5]: # The ones we know were missing
                print(f"Restored missing serial {serial} (Risk ID: {risk_id})")

        conn.commit()
        print(f"Done. Processed {restored_count} risks for the target permit.")
        conn.close()
    except Exception as e:
        print(f"Error: {e}")

if __name__ == "__main__":
    patch_missing_risks()