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

BASE_URL = "https://www.bestattung-wels.at/"
URL_TEMPLATE = f"{BASE_URL}/current-deaths?page={{}}&deathCompany={{}}"
URL_TEMPLATE_NO_COMPANY = f"{BASE_URL}/current-deaths?page={{}}"

# 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)
    dn_picture = Column(LargeBinary)
    dn_age = Column(Integer)  # New column for age

# 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",
]

class WelsCrawler:
    def __init__(self, death_company):
        self.death_company = death_company

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

    def download_picture(self, url):
        headers = {
            "User-Agent": random.choice(USER_AGENTS)
        }
        response = requests.get(url, headers=headers)
        return response.content if response.status_code == 200 else None

    def extract_data(self, soup, session):
        death_table = soup.find("table", attrs={"class": "death-table"})
        if not death_table:
            print("No death table found.")
            return False
        
        death_table_data = death_table.find_all("tr")

        for row in death_table_data:
            rows = row.find_all('td')
            cols = [ele.text.strip() for ele in rows]
            if cols:
                date_death_str = cols[1].replace('†', '').strip() if cols[1] else None
                try:
                    date_death = datetime.strptime(date_death_str, "%d.%m.%Y").date() if date_death_str else None
                except ValueError:
                    print(f"Invalid date format: {date_death_str}")
                    date_death = None
                    continue

                # full-name is available as well
                name = row.find('span', class_='full-name').text.strip()

                # name = cols[0]
                location = self.death_company.capitalize()
                
                # Extract the picture URL
                picture_tag = row.find('img', class_='profile-picture')
                picture_url = picture_tag['src'] if picture_tag else None
                picture_data = self.download_picture(picture_url) if picture_url else None

                # Extract the age
                age_str = row.find('span', class_='deceased-age').text.strip() if row.find('span', class_='deceased-age') else None
                age_match = re.search(r'(\d+)', age_str) if age_str else None
                age = int(age_match.group(1)) if age_match else None

                # Check for duplicates before inserting
                #existing_notice = session.query(DeathNotice).filter_by(dn_death=date_death, dn_name=name, dn_location=location).first()
                existing_notice = session.query(DeathNotice).filter_by(dn_death=date_death, dn_name=name).first()
                if existing_notice:
                    # update picture and age
                    existing_notice.dn_picture = picture_data
                    existing_notice.dn_age = age
                    existing_notice.dn_name = name
                    print(f"Duplicate entry found - Date: {date_death}, Name: {name}, Location: {location}")
                else:
                    # Insert data into the database using SQLAlchemy
                    death_notice = DeathNotice(dn_death=date_death, dn_name=name, dn_location=location, dn_picture=picture_data, dn_age=age)
                    session.add(death_notice)
                    print(f"Inserted data - Date: {date_death}, Name: {name}, Location: {location}, Age: {age}, Picture URL: {picture_url}")
        
        return True

    def find_next_page(self, soup, current_url):
        next_button_disabled = soup.find("a", class_="pagination__next-btn pagination__next-btn--disabled")
        if next_button_disabled:
            return None
        
        # next_button = soup.find("a", class_="pagination__next-btn")
        # if next_button:
        match = re.search(r'\?page=(\d+)', current_url) #next_button['href'])
        if match:
            next_page_num = int(match.group(1)) + 1
            if self.death_company == "unbekannt":
                next_page_url = URL_TEMPLATE_NO_COMPANY.format(next_page_num)
            else:
                next_page_url = URL_TEMPLATE.format(next_page_num, self.death_company)
            return next_page_url
        return None
    
class Worker:
    def __init__(self, crawler, start_page=1):
        self.crawler = crawler
        self.start_page = start_page

    def run(self):
        # Create a new session
        session = SessionLocal()
        
        try:
            if self.crawler.death_company == "unbekannt":
                current_url = URL_TEMPLATE_NO_COMPANY.format(self.start_page)
            else:
                current_url = URL_TEMPLATE.format(self.start_page, self.crawler.death_company)
            while current_url:
                soup = self.crawler.get_page_content(current_url)
                if not self.crawler.extract_data(soup, session):
                    break
                current_url = self.crawler.find_next_page(soup, current_url)
                if not current_url:
                    break

                # 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__":
    death_company = sys.argv[1] if len(sys.argv) > 1 else "unbekannt"
    start_page = int(sys.argv[2]) if len(sys.argv) > 2 else 1

    if death_company not in ["wels", "marchtrenk", "unbekannt"]:
        print("Invalid death company. Use 'wels' or 'marchtrenk'.")
        sys.exit(1)

    crawler = WelsCrawler(death_company)
    worker = Worker(crawler, start_page)
    worker.run()