import argparse
import logging
import os
import pathlib

import psycopg
import psycopg.sql

logging.basicConfig(level=logging.INFO, format="%(levelname)s  [%(name)s] %(message)s")
logger = logging.getLogger("fixture.main")

postgres_host = os.getenv("POSTGRES_HOST")
postgres_port = os.getenv("POSTGRES_PORT")
postgres_user = os.getenv("POSTGRES_USER")
postgres_password = os.getenv("POSTGRES_PASSWORD")
postgres_db = os.getenv("POSTGRES_DB")


def main(csv_dir: str) -> None:
    def get_connection() -> psycopg.Connection:
        return psycopg.connect(
            host=postgres_host,
            port=postgres_port,
            user=postgres_user,
            password=postgres_password,
            dbname=postgres_db,
            autocommit=True,
        )

    def get_csv_file_path(csv_dir: str) -> pathlib.Path:
        csv_path = pathlib.Path(csv_dir)
        if not csv_path.is_absolute():
            csv_path = pathlib.Path(__file__).parent / csv_path
        if not csv_path.exists():
            raise FileNotFoundError(f"CSV directory does not exist: {csv_path}")
        return csv_path

    def get_csv_files(directory: pathlib.Path) -> list[pathlib.Path]:
        return sorted(directory.glob("*.csv"), key=lambda x: x.stem)

    def get_table_name_from_csv(csv_file: pathlib.Path) -> str:
        return csv_file.stem.split(".")[1]

    def get_csv_headers(csv_file: pathlib.Path) -> list[str]:
        with pathlib.Path.open(csv_file, encoding="utf-8") as f:
            first_line = f.readline().strip()
            return first_line.split(",")

    csv_file_path = get_csv_file_path(csv_dir)
    connection = get_connection()

    # 直下のファイルを取得
    csv_files = get_csv_files(csv_file_path)
    for csv_file in csv_files:
        table_name = get_table_name_from_csv(csv_file)
        headers = get_csv_headers(csv_file)

        copy_query = psycopg.sql.SQL(
            """--sql
            COPY {} ({}) FROM STDIN WITH (FORMAT CSV, HEADER MATCH);
            """,
        ).format(
            psycopg.sql.Identifier(table_name),
            psycopg.sql.SQL(", ").join(map(psycopg.sql.Identifier, headers)),
        )
        with (
            pathlib.Path.open(csv_file, encoding="utf-8") as f,
            connection.cursor() as cursor,
            cursor.copy(copy_query) as cp,
        ):
            cp.write(f.read())
            logger.info(f"Copied data to table {table_name}.")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Insert fixture data from CSV file")
    parser.add_argument(
        "--csv-dir",
        "-d",
        type=str,
        required=True,
        default="local",
        help="Path to the CSV file containing tenant data",
    )
    args = parser.parse_args()
    main(
        csv_dir=args.csv_dir,
    )
