import csv
import datetime
import importlib
import logging
import math

from google.api_core.exceptions import NotFound
from google.cloud import bigquery

logger = logging.getLogger(__name__)


class PiExporter:
    @staticmethod
    def _clean_for_excel(text, illegal_chars_re):
        if not isinstance(text, str):
            return text
        if illegal_chars_re:
            return illegal_chars_re.sub("", text)
        return text

    @staticmethod
    def to_csv(data, filename):
        if not data:
            logger.info("No data to export.")
            return

        if not filename.endswith(".csv"):
            filename += ".csv"

        keys = data[0].keys()
        with open(filename, "w", newline="", encoding="utf-8") as handle:
            writer = csv.DictWriter(handle, fieldnames=keys)
            writer.writeheader()
            writer.writerows(data)
        logger.info("Successfully saved %s rows to %s", len(data), filename)

    @staticmethod
    def to_excel(data, filename):
        if not data:
            logger.info("No data to export.")
            return

        pandas_spec = importlib.util.find_spec("pandas")
        if pandas_spec is None:
            logger.error("Pandas is required for Excel exports.")
            return

        pd = importlib.import_module("pandas")
        from openpyxl.cell.cell import ILLEGAL_CHARACTERS_RE

        if not filename.endswith(".xlsx"):
            filename += ".xlsx"

        logger.info("Preparing Excel file: %s (%s rows)", filename, len(data))
        df = pd.DataFrame(data)
        for col in df.select_dtypes(include=["object"]).columns:
            df[col] = df[col].apply(lambda v: PiExporter._clean_for_excel(v, ILLEGAL_CHARACTERS_RE))

        df.to_excel(filename, index=False, engine="openpyxl")
        logger.info("Successfully saved %s rows to %s", len(data), filename)

    @staticmethod
    def to_google_sheet(data, spreadsheet_name, tab_name="Sheet1"):
        if not data:
            logger.info("No data to upload.")
            return

        gspread_spec = importlib.util.find_spec("gspread")
        if gspread_spec is None:
            logger.error("gspread not installed.")
            return

        gspread = importlib.import_module("gspread")
        google_auth = importlib.import_module("google.auth")
        APIError = importlib.import_module("gspread.exceptions").APIError

        row_count = len(data)
        logger.info("Preparing Google Sheet upload: %s rows", row_count)

        scopes = [
            "https://www.googleapis.com/auth/spreadsheets",
            "https://www.googleapis.com/auth/drive",
        ]
        creds, _ = google_auth.default(scopes=scopes)
        client = gspread.authorize(creds)

        try:
            sh = client.open(spreadsheet_name)
        except gspread.SpreadsheetNotFound:
            logger.info("Sheet '%s' not found. Creating it...", spreadsheet_name)
            sh = client.create(spreadsheet_name)

        try:
            worksheet = sh.worksheet(tab_name)
            worksheet.clear()
        except gspread.WorksheetNotFound:
            worksheet = sh.add_worksheet(title=tab_name, rows=row_count + 100, cols=20)

        headers = list(data[0].keys())
        rows = [[row.get(col, "") for col in headers] for row in data]

        worksheet.update([headers], "A1")

        chunk_size = 30000
        total_chunks = math.ceil(len(rows) / chunk_size)
        logger.info("Starting upload in %s chunks (size: %s)", total_chunks, chunk_size)

        for i in range(total_chunks):
            start = i * chunk_size
            end = start + chunk_size
            chunk = rows[start:end]
            try:
                worksheet.append_rows(chunk, value_input_option="RAW")
                logger.info("Uploaded chunk %s/%s (%s rows)", i + 1, total_chunks, len(chunk))
            except APIError:
                logger.exception("Google Sheets API error on chunk %s", i + 1)
                raise

        logger.info("Successfully uploaded %s rows to '%s' (tab: %s)", row_count, spreadsheet_name, tab_name)


    @staticmethod
    def _infer_bq_field(name, value):
        field_name = name.replace(" ", "_")
        if field_name == "Date":
            return bigquery.SchemaField(field_name, "DATE")
        if field_name == "Month":
            return bigquery.SchemaField(field_name, "DATE")

        if isinstance(value, bool):
            return bigquery.SchemaField(field_name, "BOOL")
        if isinstance(value, int):
            return bigquery.SchemaField(field_name, "INT64")
        if isinstance(value, float):
            return bigquery.SchemaField(field_name, "FLOAT64")
        if isinstance(value, datetime.date):
            return bigquery.SchemaField(field_name, "DATE")
        return bigquery.SchemaField(field_name, "STRING")

    @staticmethod
    def _normalize_partition_dates(data):
        for row in data:
            if "Month" in row and isinstance(row.get("Month"), str):
                month_val = row["Month"].strip()
                if month_val:
                    # Normalise YYYY-MM to a valid DATE for monthly partitioning
                    row["Month"] = f"{month_val}-01" if len(month_val) == 7 else month_val

    @staticmethod
    def _ensure_partitioned_table(client, table_ref, data):
        partition_field = None
        partition_type = None

        if "Date" in data[0]:
            partition_field = "Date"
            partition_type = bigquery.TimePartitioningType.DAY
        elif "Month" in data[0]:
            partition_field = "Month"
            partition_type = bigquery.TimePartitioningType.MONTH

        schema = [PiExporter._infer_bq_field(k, v) for k, v in data[0].items()]
        table = bigquery.Table(table_ref, schema=schema)

        if partition_field:
            table.time_partitioning = bigquery.TimePartitioning(
                type_=partition_type,
                field=partition_field,
            )

        client.create_table(table)
        logger.info("Created table %s with partitioning on %s", table_ref, partition_field or "none")

    @staticmethod
    def to_bigquery(data, project_id, dataset_id, table_id):
        if not data:
            logger.info("No data to upload.")
            return

        client = bigquery.Client(project=project_id)
        table_ref = f"{project_id}.{dataset_id}.{table_id}"

        PiExporter._normalize_partition_dates(data)

        try:
            client.get_table(table_ref)
        except NotFound:
            PiExporter._ensure_partitioned_table(client, table_ref, data)

        logger.info("Uploading %s rows to BigQuery table %s...", len(data), table_ref)

        job_config = bigquery.LoadJobConfig(
            source_format=bigquery.SourceFormat.NEWLINE_DELIMITED_JSON,
            write_disposition=bigquery.WriteDisposition.WRITE_APPEND,
        )

        job = client.load_table_from_json(data, table_ref, job_config=job_config)
        job.result()
        logger.info("Upload successful. Loaded %s rows.", job.output_rows)
