from jinja2 import Template
from openpyxl.styles import Font, PatternFill, Border, Side, Alignment
from openpyxl.utils.dataframe import dataframe_to_rows
from loguru import logger

import data
import utils
from config import Params

def process_client(country_code, client, partners):
    logger.info(f"Get data for client {client['client_name']}")
    try:
        df = data.get_country_weekly_streams(country_code, client['client_key'], partners)
        metrics = data.process_metrics(df)
        last_week_data = df.query("week_ft==@metrics['last_week']")
        main_table_columns, main_table_data = data.process_main_table(last_week_data)
        main_table_html = Template(
                            open('../html/main_table.html').read()
                            ).render(
                                columns=main_table_columns, 
                                data=main_table_data,
                                color=client['color'],
                                )
        img64 = utils.generate_img64_from_html(main_table_html)

        logger.success(f"Get data for client {client['client_name']} OK")
        return {'metrics': metrics, 'img64': img64, 'client_name': client['client_name'], 'color': client['color']}
    except Exception as e:
        logger.exception(f"Failed to get data for client {client['client_name']}: {e}")
        raise e

def create_worksheet(workbook, title, dataframe, streams_column_index):
    """
    Create and format a worksheet with the given dataframe data.
    
    Args:
        workbook: The Excel workbook object
        title: Title for the worksheet
        dataframe: DataFrame containing the data
        streams_column_index: Index of the streams column for number formatting
    
    Returns:
        The created worksheet
    """
    # Create worksheet
    if title == 'Top 50 Tracks':
        ws = workbook.active
        ws.title = title
    else:
        ws = workbook.create_sheet(title=title)
    
    # Add data to worksheet
    for r in dataframe_to_rows(dataframe, index=False, header=True):
        ws.append(r)
    
    # Define styles
    header_font = Font(bold=True)
    header_fill = PatternFill(start_color='D7E4BC', end_color='D7E4BC', fill_type='solid')
    border = Border(left=Side(style='thin'), right=Side(style='thin'), 
                    top=Side(style='thin'), bottom=Side(style='thin'))
    text_alignment = Alignment(wrap_text=True, vertical='top')
    
    # Apply header formatting
    for cell in ws[1]:
        cell.font = header_font
        cell.fill = header_fill
        cell.border = border
        cell.alignment = text_alignment
    
    # Apply data formatting and borders
    for row in ws.iter_rows(min_row=2, max_row=ws.max_row, min_col=1, max_col=ws.max_column):
        for cell in row:
            cell.border = border
            if cell.column in [1, streams_column_index]:  # Rank and Streams columns
                cell.number_format = '#,##0'
            else:
                cell.alignment = text_alignment
    
    # Auto-adjust column widths
    for column in ws.columns:
        max_length = 0
        column_letter = column[0].column_letter
        for cell in column:
            try:
                if len(str(cell.value)) > max_length:
                    max_length = len(str(cell.value))
            except:
                pass
        adjusted_width = min(max_length + 20, 50)
        ws.column_dimensions[column_letter].width = adjusted_width
    
    return ws