from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker
import json
import logging

# Setup DB
DB_PATH = "sqlite:///companies_v3_fixed_2.db"
engine = create_engine(DB_PATH)
SessionLocal = sessionmaker(bind=engine)
session = SessionLocal()

# Import Models (Simplified for script)
from sqlalchemy import Column, Integer, String, Text, JSON
from sqlalchemy.ext.declarative import declarative_base

Base = declarative_base()

class Company(Base):
    __tablename__ = "companies"
    id = Column(Integer, primary_key=True)
    name = Column(String)
    city = Column(String)
    country = Column(String)
    crm_vat = Column(String)
    street = Column(String)
    zip_code = Column(String)

class EnrichmentData(Base):
    __tablename__ = "enrichment_data"
    id = Column(Integer, primary_key=True)
    company_id = Column(Integer)
    source_type = Column(String)
    content = Column(JSON)

def fix_data():
    company_id = 32
    print(f"🔧 Fixing Data for Company ID {company_id}...")
    
    company = session.query(Company).filter_by(id=company_id).first()
    if not company:
        print("❌ Company not found.")
        return

    enrichment = session.query(EnrichmentData).filter_by(
        company_id=company_id, source_type="website_scrape"
    ).first()
    
    if enrichment and enrichment.content:
        imp = enrichment.content.get("impressum")
        if imp:
            print(f"📄 Found Impressum: {imp}")
            
            changed = False
            if imp.get("city"):
                company.city = imp.get("city")
                changed = True
                print(f"  -> Set City: {company.city}")
                
            if imp.get("vat_id"):
                company.crm_vat = imp.get("vat_id")
                changed = True
                print(f"  -> Set VAT: {company.crm_vat}")
            
            if imp.get("country_code"):
                company.country = imp.get("country_code")
                changed = True
                print(f"  -> Set Country: {company.country}")

            if imp.get("street"):
                company.street = imp.get("street")
                changed = True
                print(f"  -> Set Street: {company.street}")

            if imp.get("zip"):
                company.zip_code = imp.get("zip")
                changed = True
                print(f"  -> Set Zip: {company.zip_code}")

            if changed:
                session.commit()
                print("✅ Database updated.")
            else:
                print("ℹ️ No changes needed.")
        else:
            print("⚠️ No impressum data in enrichment.")
    else:
        print("⚠️ No enrichment data found.")

if __name__ == "__main__":
    fix_data()