import os
import io
import time
import pandas as pd
from lawrisk.services.licensing_repo import start_permit_import_session, commit_permit_import_session, _lic_pg_conn
from lawrisk.utils.env_loader import load_env

def test_import_searchability():
    load_env()
    print("Starting test import...")
    
    # Create simple Excel in memory
    # Using columns that are known to work
    data = [
        ["事项名称", "风险提示内容", "法律依据", "主题"],
        ["测试许可_SearchTest", "测试风险内容", "测试依据", "测试主题_SearchTest"]
    ]
    df = pd.DataFrame(data[1:], columns=data[0])
    
    output = io.BytesIO()
    with pd.ExcelWriter(output, engine='openpyxl') as writer:
        df.to_excel(writer, index=False, sheet_name="测试地区_SearchTest")
    excel_data = output.getvalue()
    
    try:
        # 1. Start session
        session = start_permit_import_session(
            file_bytes=excel_data,
            filename="test_search.xlsx",
            binding_mode="auto"
        )
        session_id = session['session_id']
        sheet_summaries = session['sheet_summaries']
        print(f"Session started: {session_id}")
        for s in sheet_summaries:
            print(f" - Sheet: {s['sheet_name']}, row_count: {s['row_count']}, permit_count: {s['permit_count']}")
            if s['missing_region']:
                print(f"   WARNING: missing_region for {s['sheet_name']}")
        
        sheet_names = [s['sheet_name'] for s in sheet_summaries]
        
        # 2. Commit session
        result = commit_permit_import_session(
            session_id=session_id,
            sheet_names=sheet_names,
            theme_bindings={
                 "测试地区_SearchTest": {
                     "测试许可_SearchTest": ["测试主题_SearchTest"]
                 }
            },
            edited_by="test_bot"
        )
        print("Import committed successfully.")
        
        # 3. Check database for region_themes
        with _lic_pg_conn() as conn:
            cur = conn.cursor()
            # Wait a tiny bit just in case of async issues which shouldn't happen here
            cur.execute("""
                SELECT r.name, t.name 
                FROM region_themes rt
                JOIN regions r ON r.id = rt.region_id
                JOIN themes t ON t.id = rt.theme_id
                WHERE r.name = '测试地区_SearchTest'
            """)
            rows = cur.fetchall()
            print(f"Entries in region_themes for test region: {rows}")
            
            if rows:
                print("SUCCESS: region_themes entry was created during import.")
            else:
                # Let's check region_theme_permits to see if anything was imported at all
                cur.execute("""
                    SELECT r.name, p.name, t.name
                    FROM region_theme_permits rtp
                    JOIN regions r ON r.id = rtp.region_id
                    JOIN permits p ON p.id = rtp.permit_id
                    LEFT JOIN themes t ON t.id = rtp.theme_id
                    WHERE r.name = '测试地区_SearchTest'
                """)
                rtp_rows = cur.fetchall()
                print(f"Entries in region_theme_permits: {rtp_rows}")
                print("FAILURE: region_themes entry was NOT created during import.")
                
            # Clean up test data
            print("Cleaning up test data...")
            cur.execute("SELECT id FROM regions WHERE name = '测试地区_SearchTest'")
            res = cur.fetchone()
            if res:
                rid = res[0]
                cur.execute("DELETE FROM region_theme_permits WHERE region_id = %s", (rid,))
                cur.execute("DELETE FROM region_permit_details WHERE region_id = %s", (rid,))
                cur.execute("DELETE FROM region_permit_risks WHERE region_id = %s", (rid,))
                cur.execute("DELETE FROM regions WHERE id = %s", (rid,))
            
            conn.commit()
            print("Cleanup done.")

    except Exception as e:
        print(f"Test failed with error: {e}")
        # import traceback
        # traceback.print_exc()

if __name__ == "__main__":
    test_import_searchability()