import csv
from abc import ABC, abstractmethod
from os.path import join
from typing import Any, List

import boto3
import pandas as pd
import pymysql

from schema import Table


class DataBackend(ABC):
    @abstractmethod
    def rebuild_schemas(self, tables: List[Table]):
        pass

    @abstractmethod
    def insert_data(self, table: Table, values: List[Any]):
        pass


class TiDBBackend(DataBackend):
    def __init__(self, connection: pymysql.Connection):
        super().__init__()
        self.connection = connection

    def rebuild_schemas(self, tables):
        for table in tables:
            self.prepare_database(table.schema_name())
            self.prepare_table(table.show_create_table())

    def insert_data(self, table: Table, values: List[List[Any]]):
        column_names = [column.name for column in table.columns]
        sql_template = f"""
            replace into `{table.schema_name()}`.`{table.name}` ({",".join(column_names)})
            values ({",".join(['%s' for _ in range(len(column_names))])})
        """

        with self.connection.cursor() as cur:
            cur.executemany(sql_template, values)
        self.connection.commit()

    def prepare_database(self, database: str):
        with self.connection.cursor() as cur:
            cur.execute(f"CREATE DATABASE IF NOT EXISTS `{database}`")

    def prepare_table(self, ddl: str):
        with self.connection.cursor() as cur:
            cur.execute(ddl)


class S3Backend(DataBackend):
    def __init__(self, bucket: str, base_dir: str):
        super().__init__()
        self.bucket = bucket
        self.base_dir = base_dir
        self.created_schemas = set()
        self.table_indexes = {}

    def rebuild_schemas(self, tables: List[Table]):
        client = boto3.client("s3")
        for table in tables:
            if table.schema_name() not in self.created_schemas:
                client.put_object(
                    Bucket=self.bucket,
                    Body=f"CREATE DATABASE IF NOT EXISTS `{table.schema_name()}`;".encode(),
                    Key=join(self.base_dir, f"{table.schema_name()}-schema-create.sql"),
                )
                self.created_schemas.add(table.schema_name())

            ddl = table.show_create_table() + ";"
            client.put_object(
                Bucket=self.bucket,
                Body=ddl.encode(),
                Key=join(
                    self.base_dir, f"{table.schema_name()}.{table.name}-schema.sql"
                ),
            )

    def insert_data(self, table: Table, values: List[List[Any]]):
        client = boto3.client("s3")
        key = f"{table.schema_name()}.{table.name}"
        if key not in self.table_indexes:
            self.table_indexes[key] = 0

        df = pd.DataFrame(values, columns=[column.name for column in table.columns])
        client.put_object(
            Bucket=self.bucket,
            Body=df.to_csv(
                index=False, escapechar="\\", quoting=csv.QUOTE_ALL
            ).encode(),
            Key=join(self.base_dir, f"{key}.{self.table_indexes[key]}.csv"),
        )
        self.table_indexes[key] += 1