import os
import sys
import shutil
import logging
from datetime import date
import requests
import markdown
from jinja2 import Environment, FileSystemLoader

# Assuming these are in the same directory or Python path
from import_helpers import music_data, ExtensionAlbum, ExtensionSong, ExtensionReview
from audio_downloader import process_html_files
from markdown_extensions import AudioExtension

# Constants specific to music generation
CURRENT_DIR = os.path.dirname(os.path.abspath(__file__))
SOURCE_ASSETS_DIR = os.path.join(CURRENT_DIR, "assets")
SOURCE_COVER_DIR = os.path.join(SOURCE_ASSETS_DIR, "album_covers")
SOURCE_AUDIO_DIR = os.path.join(SOURCE_ASSETS_DIR, "audio")
# The relative path used in HTML *remains* the same, pointing within assets
OUTPUT_AUDIO_SUBDIR = "/assets/audio"


TARGET_ALBUM_IDS = [
    "OLAK5uy_nLmxB70YJBLEThZljF4v1xhhSJ2pwB6GA", # spelling
    "OLAK5uy_lvJG6wOQoiTe_vyGkre-RbZWnd4y5NtXM", # congratulations
    "OLAK5uy_kMVviQYyPZN-EUsk91z_hZJzCzu1804xg", # imaginal disk
    "OLAK5uy_kmzoSOa_tCizE-r4sweNz91d9qBv1UCVY", # brat
    "OLAK5uy_mgDUvSi1yFc9Sr-s1BmCGZKBFGravXSEc", # sinister grift
    "OLAK5uy_mlZk7kmAEKb8ZQWs-2sxnMmRrANtpXpzU", # 13"
    "OLAK5uy_mXoi-FuQb9Gw7Mguhdx5F4jltT0L1qOCw", # dsotm
    "OLAK5uy_mPI0_Yf-l12oLr5cbmu_HorkcRRDYPMfE", # vespertine
    # Add more album IDs here if needed
]

MUSIC_INDEX_MD_PATH = os.path.join(CURRENT_DIR, "md/music_index.md") 
# --- Helper Functions ---

def format_date(date_obj: date) -> str:
    """Formats a date object into 'Mon DD \'YY' format."""
    if not date_obj:
        return ""
    return date_obj.strftime("%b '%y")

def download_album_cover(album_id: str, album: ExtensionAlbum, cover_dir: str) -> str | None:
    """Downloads album cover if not already present, returns the local path."""
    if not album or not album.thumbnail_url:
        logging.warning(f"Album {album_id} has no thumbnail URL. Skipping download.")
        return None

    local_image_filename = f"{album_id}.jpg"
    local_image_path = os.path.join(cover_dir, local_image_filename)

    if os.path.exists(local_image_path):
        logging.debug(f"Cover already exists for {album_id}: {local_image_path}")
        return local_image_path

    logging.info(f"Downloading cover for {album_id} from {album.thumbnail_url}...")
    try:
        response = requests.get(album.thumbnail_url, stream=True, timeout=10)
        response.raise_for_status()

        with open(local_image_path, 'wb') as f:
            for chunk in response.iter_content(chunk_size=8192):
                f.write(chunk)
        logging.info(f"Successfully downloaded cover for {album_id} to {local_image_path}")
        return local_image_path

    except requests.exceptions.RequestException as e:
        logging.error(f"Error downloading image {album.thumbnail_url} for album {album_id}: {e}")
        if os.path.exists(local_image_path):
            os.remove(local_image_path)
        return None
    except Exception as e:
        logging.error(f"An unexpected error occurred during download for {album_id}: {e}")
        if os.path.exists(local_image_path):
            os.remove(local_image_path)
        return None

def render_album_html(album: ExtensionAlbum, artist_str: str, album_reviews_data: list,
                      songs_with_reviews_data: list, output_dir: str, env: Environment, cover_path: str | None) -> dict | None:
    """Renders the HTML for a single album using the template."""
    try:
        template = env.get_template('templates/album_template.html')
    except Exception as e:
        logging.error(f"Error loading album_template.html: {e}")
        return None

    html_content = template.render(
        album=album,
        artist_str=artist_str,
        album_reviews=album_reviews_data,
        songs_with_reviews=songs_with_reviews_data,
        cover_path=cover_path
    )

    filename = f"{album.id}.html"
    output_path = os.path.join(output_dir, filename)
    try:
        with open(output_path, "w", encoding="utf-8") as f:
            f.write(html_content)
        return {
            "filename": filename,
            "album_name": album.name,
            "artist_str": artist_str
            # Cover path is added later after checking download
        }
    except Exception as e:
        logging.error(f"Error writing HTML to {output_path} for album {album.id}: {e}")
        return None

def generate_index_html(pages_info: list, site_description_html: str, output_dir: str, env: Environment):
    """Generates the index.html file linking to all album pages."""
    # Load the main generic list page template
    template = env.get_template('templates/list_page_template.html')
    # Load the template for individual list items
    item_template = env.get_template('templates/list_templates/music_album_item.html')


    # Render each list item individually
    list_items_html = []
    for page_data in pages_info:
        # Ensure cover_path exists, provide placeholder if not (handled in item template now)
        # page_data['cover_path'] = page_data.get('cover_path') # Or set placeholder here
        list_items_html.append(item_template.render(page=page_data))

    # Render the main index page template with the list of item HTML
    html_content = template.render(
        page_title="WU - Music", # Set a specific title for this page
        site_description_html=site_description_html,
        list_items_html=list_items_html, # Pass the list of rendered HTML strings
        should_include_audio_player=True # Explicitly include audio player for music index
    )

    output_path = os.path.join(output_dir, "index.html")
    try:
        with open(output_path, "w", encoding="utf-8") as f:
            f.write(html_content)
        logging.debug(f"Successfully generated music index.html -> {output_path}")
    except Exception as e:
        logging.error(f"Error writing music index.html to {output_path}: {e}")

# --- Main Generation Function ---

def generate_music_site(output_dir: str, templates_dir: str):
    """Generates the complete music website section."""

    logging.info(f"Starting music site generation process for {len(TARGET_ALBUM_IDS)} target albums...")
    logging.info(f"Output directory set to: {output_dir}")

    # Create output directory if it doesn't exist
    os.makedirs(output_dir, exist_ok=True)
    # Create source asset directories needed before processing
    os.makedirs(SOURCE_COVER_DIR, exist_ok=True)
    os.makedirs(SOURCE_AUDIO_DIR, exist_ok=True) # Ensure this exists for downloads

    # Set up Jinja2 environment using the provided templates directory
    jinja_env = Environment(loader=FileSystemLoader(templates_dir))

    # --- Data Storage ---
    albums_data = {}
    image_paths = {}

    for album_id in TARGET_ALBUM_IDS:
        album: ExtensionAlbum = music_data.all_albums.get(album_id)
        if album:
            albums_data[album_id] = album
        else:
            logging.warning(f"Album data not found for ID: {album_id}. Skipping.")

    active_album_ids = list(albums_data.keys())
    logging.info(f"Found data for {len(active_album_ids)} albums to process.")

    # --- Step 1: Download Album Covers (if needed) ---
    logging.info("--- Music Step 1: Checking/Downloading Album Covers ---")
    for album_id in active_album_ids:
        album = albums_data[album_id]
        # Download to the *source* cover directory
        image_path = download_album_cover(album_id, album, SOURCE_COVER_DIR)
        image_paths[album_id] = image_path # Store path (or None if failed)

    # --- Step 2: Generate HTML for each album ---
    logging.info("--- Music Step 2: Generating Album HTML Pages ---")
    generated_pages_info = []
    for album_id in active_album_ids:
        album = albums_data[album_id]

        artist_names = [music_data.artists.get(artist_id).name
                        for artist_id in album.artists
                        if music_data.artists.get(artist_id)]
        artist_str = ", ".join(artist_names) if artist_names else "Unknown Artist"

        album_reviews_data = []
        album_md_parser = markdown.Markdown(extensions=[AudioExtension(fallback_song_id=None, relative_audio_path=OUTPUT_AUDIO_SUBDIR)])
        for review in album.reviews:
            html_content = album_md_parser.convert(review.markdown)
            album_reviews_data.append({
                "html_text": html_content,
                "formatted_date": format_date(review.addedOn)
            })
            album_md_parser.reset()

        songs_with_reviews_data = []
        for song_id in album.songIds:
            song: ExtensionSong = music_data.songs.get(song_id)
            if song and song.reviews:
                song_reviews_data = []
                song_md_parser = markdown.Markdown(extensions=[AudioExtension(fallback_song_id=song_id, relative_audio_path=OUTPUT_AUDIO_SUBDIR)])
                for review in song.reviews:
                    html_content = song_md_parser.convert(review.markdown)
                    song_reviews_data.append({
                        "html_text": html_content,
                        "formatted_date": format_date(review.addedOn)
                    })
                    song_md_parser.reset()
                if song_reviews_data:
                    songs_with_reviews_data.append({
                        "id": song.id,
                        "name": song.name,
                        "reviews": song_reviews_data
                    })

        local_cover_path_relative_to_output = None
        full_image_path = image_paths.get(album_id)
        if full_image_path:
            image_filename = os.path.basename(full_image_path)
            # Path relative to the *output* directory where the HTML file will be
            # Assumes 'assets/album_covers' structure within the output dir
            local_cover_path_relative_to_output = os.path.join("/assets", "album_covers", image_filename)

        page_info = render_album_html(
            album=album,
            artist_str=artist_str,
            album_reviews_data=album_reviews_data,
            songs_with_reviews_data=songs_with_reviews_data,
            output_dir=output_dir, # Use the passed output_dir
            env=jinja_env,
            cover_path=local_cover_path_relative_to_output # Pass relative path for template
        )

        if page_info:
            # Add the relative cover path needed for the index page rendering
            page_info["cover_path"] = local_cover_path_relative_to_output
            generated_pages_info.append(page_info)
            logging.debug(f"Successfully generated HTML for {album.name}")
        else:
             logging.error(f"Failed to generate HTML for {album.name} ({album_id})")

    # --- Step 3: Generate Music Index HTML ---
    logging.info("--- Music Step 3: Generating Music Index HTML ---")

    site_description_html = f"<p>Error: {MUSIC_INDEX_MD_PATH} not found or could not be read.</p>"
    if os.path.exists(MUSIC_INDEX_MD_PATH):
        try:
            with open(MUSIC_INDEX_MD_PATH, "r", encoding="utf-8") as f:
                index_md_content = f.read()
            index_md_parser = markdown.Markdown(extensions=[AudioExtension(fallback_song_id=None, relative_audio_path=OUTPUT_AUDIO_SUBDIR)])
            site_description_html = index_md_parser.convert(index_md_content)
            logging.info(f"Successfully read and converted {MUSIC_INDEX_MD_PATH}")
        except Exception as e:
            logging.error(f"Error processing {MUSIC_INDEX_MD_PATH}: {e}")
    else:
         logging.warning(f"{MUSIC_INDEX_MD_PATH} not found. Using default description.")

    if generated_pages_info:
        generate_index_html(generated_pages_info, site_description_html, output_dir, jinja_env)
    else:
        logging.warning("No album pages were generated successfully for music. Index file not created.")

    # --- Step 4: Download Audio Segments ---
    logging.info("--- Music Step 4: Processing and Downloading Audio Segments ---")
    # process_html_files needs the directory containing the HTML files (output_dir)
    # and the *source* directory where audio should be downloaded.
    process_html_files(html_dir=output_dir, audio_output_dir=SOURCE_AUDIO_DIR)

    logging.info("Music site generation process finished.")


# Example of how to call this from another script:
# if __name__ == "__main__":
#     # Configure logging if running standalone
#     log_format = '%(asctime)s - %(name)s - %(levelname)s: %(message)s'
#     logging.basicConfig(level=logging.INFO, format=log_format, stream=sys.stdout)
#     logging.getLogger().setLevel(logging.INFO)
#     logging.getLogger("urllib3").setLevel(logging.WARNING)

#     # Define paths relative to the *calling* script or CWD
#     base_output_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "output")
#     music_output_dir = os.path.join(base_output_dir, "music")
#     templates_dir = os.path.dirname(os.path.abspath(__file__)) # Assuming templates are here
#     music_index_md_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "index.md") # Assuming index.md is here

#     generate_music_site(music_output_dir, templates_dir, music_index_md_path) 