import click
import pandas as pd
from lxml import etree
import gzip
import os
from tqdm import tqdm
from datetime import datetime

# Generate a sitemap XML string
def generate_sitemap(urls, priority=None, frequency=None, lastmod=None):
    urlset = etree.Element('urlset', xmlns="http://www.sitemaps.org/schemas/sitemap/0.9")
    comment = etree.Comment("Generated by seo-bhishma-cli tool")
    urlset.append(comment)
    for url in urls:
        url_element = etree.Element('url')
        loc_element = etree.Element('loc')
        loc_element.text = url
        url_element.append(loc_element)
        if lastmod:
            lastmod_element = etree.Element('lastmod')
            lastmod_element.text = lastmod
            url_element.append(lastmod_element)
        if priority:
            priority_element = etree.Element('priority')
            priority_element.text = priority
            url_element.append(priority_element)
        if frequency:
            frequency_element = etree.Element('changefreq')
            frequency_element.text = frequency
            url_element.append(frequency_element)
        urlset.append(url_element)
    return etree.tostring(urlset, pretty_print=True, xml_declaration=True, encoding='UTF-8')

# Write sitemap to file
def write_sitemap(file_path, sitemap_content, compressed=False):
    if compressed:
        with gzip.open(file_path, 'wb') as f:
            f.write(sitemap_content)
    else:
        with open(file_path, 'wb') as f:
            f.write(sitemap_content)

# Read input CSV file
def read_input_file(file_path):
    return pd.read_csv(file_path)['url'].tolist()

@click.command()
@click.pass_context
def sitemap_generator(ctx):
    """Generate XML sitemaps from a list of URLs."""
    while True:
        click.echo("\n" + "="*50)
        click.echo(click.style("Sitemap Generator", fg="yellow", bold=True))
        click.echo(click.style("1. Generate a single sitemap", fg="cyan"))
        click.echo(click.style("2. Generate nested sitemaps", fg="cyan"))
        click.echo(click.style("3. Exit", fg="red", bold=True))
        choice = click.prompt(click.style("Enter your choice", fg="cyan", bold=True), type=int)

        if choice == 3:
            click.echo(click.style("Exiting Sitemap Generator. Goodbye!", fg="red", bold=True))
            break
        elif choice in [1, 2]:
            input_file = click.prompt(click.style("Enter the path to the input CSV file", fg="cyan"), type=click.Path(exists=True))
            output_dir = click.prompt(click.style("Enter the output directory", fg="cyan"), type=click.Path())
            
            if not os.path.exists(output_dir):
                click.echo(click.style(f"Output directory '{output_dir}' does not exist.", fg="red"))
                click.echo(click.style("Please create the directory and try again.", fg="red"))
                continue

            nested = choice == 2
            url_limit = 50000
            if nested:
                url_limit = click.prompt(click.style("Enter the maximum number of URLs per sitemap", fg="cyan"), default=50000, show_default=True, type=int)
            compressed = click.confirm(click.style("Do you want to create compressed sitemaps?", fg="cyan"))
            priority = click.prompt(click.style("Enter priority for URLs (or leave blank)", fg="cyan"), default='', show_default=False)
            frequency = click.prompt(click.style("Enter change frequency for URLs (or leave blank)", fg="cyan"), default='', show_default=False)
            lastmod = click.prompt(click.style("Enter the last modified date for URLs (or leave blank for current time)", fg="cyan"), default='', show_default=False)
            if not lastmod:
                lastmod = datetime.now().strftime('%Y-%m-%d')

            urls = read_input_file(input_file)
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")

            if not nested:
                # Single sitemap
                click.echo(click.style("Generating single sitemap...", fg="green", bold=True))
                sitemap_content = generate_sitemap(urls, priority if priority else None, frequency if frequency else None, lastmod)
                file_name = f'sitemap_{timestamp}.xml'
                file_path = os.path.join(output_dir, file_name)
                if compressed:
                    file_path += '.gz'
                write_sitemap(file_path, sitemap_content, compressed)
                click.echo(click.style(f"Single sitemap saved to {file_path}", fg="green", bold=True))
            else:
                # Nested sitemaps
                click.echo(click.style("Generating nested sitemaps...", fg="green", bold=True))
                sitemap_index = []
                for i in tqdm(range(0, len(urls), url_limit), desc="Creating sitemaps"):
                    sitemap_urls = urls[i:i+url_limit]
                    sitemap_content = generate_sitemap(sitemap_urls, priority if priority else None, frequency if frequency else None, lastmod)
                    sitemap_file_name = f'sitemap_{timestamp}_{i // url_limit}.xml'
                    sitemap_file_path = os.path.join(output_dir, sitemap_file_name)
                    if compressed:
                        sitemap_file_path += '.gz'
                    write_sitemap(sitemap_file_path, sitemap_content, compressed)
                    sitemap_index.append(sitemap_file_name if not compressed else sitemap_file_name + '.gz')
                
                # Generate sitemap index
                sitemapindex = etree.Element('sitemapindex', xmlns="http://www.sitemaps.org/schemas/sitemap/0.9")
                comment = etree.Comment("Generated by seo-bhishma-cli tool")
                sitemapindex.append(comment)
                for sitemap in sitemap_index:
                    sitemap_element = etree.Element('sitemap')
                    loc_element = etree.Element('loc')
                    loc_element.text = os.path.join(output_dir, sitemap)
                    sitemap_element.append(loc_element)
                    sitemapindex.append(sitemap_element)
                
                sitemap_index_content = etree.tostring(sitemapindex, pretty_print=True, xml_declaration=True, encoding='UTF-8')
                sitemap_index_file = f'sitemap_index_{timestamp}.xml'
                sitemap_index_path = os.path.join(output_dir, sitemap_index_file)
                if compressed:
                    sitemap_index_path += '.gz'
                write_sitemap(sitemap_index_path, sitemap_index_content, compressed)
                click.echo(click.style(f"Sitemap index saved to {sitemap_index_path}", fg="green", bold=True))
        else:
            click.echo(click.style("Invalid choice. Please select a valid option.", fg="red"))

        click.echo("\n" + "="*50 + "\n")

if __name__ == '__main__':
    sitemap_generator()
