import requests
from bs4 import BeautifulSoup
import re
import yaml
from sqlalchemy import create_engine, Column, Integer, String, Date, UniqueConstraint
from sqlalchemy.orm import declarative_base, sessionmaker
from datetime import datetime
import sys
import time
import random

BASE_URL = "https://www.bestattung-aichinger.at"
URL_TEMPLATE = f"{BASE_URL}/traueranzeigen--5683197-de.html?p={{}}"

# Load database credentials from config.yaml
with open("config.yaml", "r") as file:
    config = yaml.safe_load(file)

db_config = config['database']
DATABASE_URL = f"postgresql://{db_config['username']}:{db_config['password']}@{db_config['host']}:{db_config['port']}/{db_config['dbname']}"

# Define the SQLAlchemy base and model
Base = declarative_base()

class DeathNotice(Base):
    __tablename__ = 'death_notices'
    __table_args__ = (
        UniqueConstraint('dn_death', 'dn_name', 'dn_location', name='unique_death_notice'),
        {'schema': 'pomos'}
    )
    id = Column(Integer, primary_key=True, autoincrement=True)
    dn_death = Column(Date)
    dn_name = Column(String)
    dn_location = Column(String)

# Create the SQLAlchemy engine and session
engine = create_engine(DATABASE_URL)
SessionLocal = sessionmaker(bind=engine)

# Create the schema and tables
Base.metadata.create_all(engine)

# List of User-Agent strings to rotate
USER_AGENTS = [
    "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/58.0.3029.110 Safari/537.3",
    "Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:54.0) Gecko/20100101 Firefox/54.0",
    "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Edge/16.16299",
    "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_13_4) AppleWebKit/605.1.15 (KHTML, like Gecko) Version/11.1 Safari/605.1.15",
    "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/57.0.2987.110 Safari/537.36",
]

def get_page_content(url):
    headers = {
        "User-Agent": random.choice(USER_AGENTS)
    }
    response = requests.get(url, headers=headers)
    return BeautifulSoup(response.content, "html.parser")

def extract_data(soup, session):
    death_table = soup.find("ul", attrs={"class": "homepage_unterseiten_layout_todesanzeigen"})
    if not death_table:
        print("No death table found.")
        return False
    
    death_table_data = death_table.find_all("li")
    for row in death_table_data:
        date_span = row.find('span', class_='homepage_unterseiten_layout_datum')
        date_death_str = date_span.text.replace('†', '').strip() if date_span else None

        # Parse the date string to a datetime.date object
        date_death = datetime.strptime(date_death_str, "%d.%m.%Y").date() if date_death_str else None

        title_span = row.find('span', class_='homepage_unterseiten_layout_titel')

        if title_span:
            title_text = title_span.get_text(separator='\n')
            title_parts = title_text.split('\n')
            title_left = title_parts[0].strip() if len(title_parts) > 0 else None
            title_right = title_parts[1].strip() if len(title_parts) > 1 else None
        else:
            title_left = title_right = None

        print(f"Extracted data - Date: {date_death}, Name: {title_left}, Location: {title_right}")
        
        # Check for duplicates before inserting
        existing_notice = session.query(DeathNotice).filter_by(dn_death=date_death, dn_name=title_left, dn_location=title_right).first()
        if existing_notice:
            print(f"Duplicate entry found - Date: {date_death}, Name: {title_left}, Location: {title_right}")
        else:
            # Insert data into the database using SQLAlchemy
            death_notice = DeathNotice(dn_death=date_death, dn_name=title_left, dn_location=title_right)
            session.add(death_notice)
            print(f"Inserted data - Date: {date_death}, Name: {title_left}, Location: {title_right}")
    
    return True

def find_next_page(current_url):
    match = re.search(r'\?p=(\d+)', current_url)
    if match:
        next_page_num = int(match.group(1)) + 1
    else:
        next_page_num = 2
    next_page_url = f"{BASE_URL}/traueranzeigen--5683197-de.html?p={next_page_num}"
    return next_page_url

def main(start_page=1):
    # Create a new session
    session = SessionLocal()
    
    try:
        current_url = URL_TEMPLATE.format(start_page)
        while current_url:
            soup = get_page_content(current_url)
            if not extract_data(soup, session):
                break
            current_url = find_next_page(current_url)

            # Random delay between requests
            time.sleep(random.uniform(1, 5))

        # Commit the transaction after processing all rows on the page
        session.commit()
    except Exception as e:
        print(f"An error occurred: {e}")
        session.rollback()
    finally:
        # Close the session
        session.close()

if __name__ == "__main__":
    start_page = int(sys.argv[1]) if len(sys.argv) > 1 else 1
    main(start_page)