from __future__ import annotations

import argparse
import os
import re
import sys
import time
from datetime import datetime
from pathlib import Path

import pymysql


MARKET_SUFFIX = {
    "SH": ".SH",
    "SZ": ".SZ",
    "BJ": ".BJ",
}


def connect(args: argparse.Namespace):
    password = args.password or os.environ.get("TIANXIA_MYSQL_PASSWORD")
    if not password:
        raise RuntimeError("missing MySQL password: pass --password or set TIANXIA_MYSQL_PASSWORD")
    return pymysql.connect(
        host=args.host,
        port=args.port,
        user=args.user,
        password=password,
        database=args.database,
        charset="utf8mb4",
        autocommit=False,
        local_infile=True,
    )


def ensure_schema(conn) -> None:
    ddl = [
        """
        CREATE TABLE IF NOT EXISTS a_share_minute_price_full_adj (
          trade_date date NOT NULL,
          trade_time time NOT NULL,
          bar_time datetime NOT NULL,
          symbol char(9) NOT NULL,
          market varchar(8) NOT NULL,
          open_price decimal(14,4) NULL,
          high_price decimal(14,4) NULL,
          low_price decimal(14,4) NULL,
          close_price decimal(14,4) NULL,
          volume decimal(24,4) NULL,
          amount decimal(24,4) NULL,
          source_file varchar(255) NOT NULL,
          source_archive varchar(255) NOT NULL,
          import_run_id varchar(96) NOT NULL,
          imported_at timestamp NOT NULL DEFAULT CURRENT_TIMESTAMP,
          PRIMARY KEY (symbol, trade_date, trade_time),
          KEY idx_full_adj_minute_date_symbol_time (trade_date, symbol, trade_time),
          KEY idx_full_adj_minute_run (import_run_id)
        ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4
        """,
        """
        CREATE TABLE IF NOT EXISTS a_share_minute_full_adj_import_run (
          run_id varchar(96) NOT NULL,
          started_at datetime NOT NULL,
          finished_at datetime NULL,
          status varchar(32) NOT NULL,
          source_archive varchar(512) NOT NULL,
          source_root varchar(512) NOT NULL,
          files_total int NOT NULL DEFAULT 0,
          files_loaded int NOT NULL DEFAULT 0,
          rows_loaded bigint NOT NULL DEFAULT 0,
          min_trade_date date NULL,
          max_trade_date date NULL,
          error_message text NULL,
          PRIMARY KEY (run_id)
        ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4
        """,
        """
        CREATE TABLE IF NOT EXISTS a_share_minute_full_adj_import_file_log (
          id bigint NOT NULL AUTO_INCREMENT,
          run_id varchar(96) NOT NULL,
          symbol char(9) NOT NULL,
          market varchar(8) NOT NULL,
          source_file varchar(512) NOT NULL,
          file_size bigint NOT NULL,
          rows_loaded bigint NOT NULL,
          started_at datetime NOT NULL,
          finished_at datetime NOT NULL,
          duration_ms bigint NOT NULL,
          status varchar(32) NOT NULL,
          error_message text NULL,
          PRIMARY KEY (id),
          UNIQUE KEY uq_full_adj_file_run (run_id, source_file),
          KEY idx_full_adj_file_symbol (symbol),
          KEY idx_full_adj_file_status (status)
        ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4
        """,
    ]
    with conn.cursor() as cur:
        for sql in ddl:
            cur.execute(sql)
    conn.commit()


def discover_files(root: Path) -> list[Path]:
    files = sorted(root.rglob("price_*.csv"))
    if not files:
        raise RuntimeError(f"no price_*.csv files found under {root}")
    return files


def symbol_for(path: Path) -> tuple[str, str]:
    market = path.parent.name.upper()
    suffix = MARKET_SUFFIX.get(market)
    if not suffix:
        raise RuntimeError(f"unsupported market directory {market}: {path}")
    m = re.match(r"price_(\d{6})\.csv$", path.name, re.IGNORECASE)
    if not m:
        raise RuntimeError(f"unsupported file name: {path}")
    return m.group(1) + suffix, market


def init_run(conn, run_id: str, source_archive: str, source_root: str, files_total: int) -> None:
    with conn.cursor() as cur:
        cur.execute(
            """
            INSERT INTO a_share_minute_full_adj_import_run
              (run_id, started_at, status, source_archive, source_root, files_total)
            VALUES (%s, NOW(), 'running', %s, %s, %s)
            ON DUPLICATE KEY UPDATE
              status='running',
              source_archive=VALUES(source_archive),
              source_root=VALUES(source_root),
              files_total=VALUES(files_total),
              error_message=NULL
            """,
            (run_id, source_archive, source_root, files_total),
        )
    conn.commit()


def finish_run(conn, run_id: str, status: str, error: str | None = None) -> None:
    with conn.cursor() as cur:
        cur.execute(
            """
            UPDATE a_share_minute_full_adj_import_run r
            LEFT JOIN (
              SELECT
                import_run_id,
                COUNT(*) AS rows_loaded,
                MIN(trade_date) AS min_trade_date,
                MAX(trade_date) AS max_trade_date
              FROM a_share_minute_price_full_adj
              WHERE import_run_id=%s
              GROUP BY import_run_id
            ) x ON x.import_run_id = r.run_id
            SET
              r.finished_at=NOW(),
              r.status=%s,
              r.rows_loaded=COALESCE(x.rows_loaded, r.rows_loaded),
              r.min_trade_date=COALESCE(x.min_trade_date, r.min_trade_date),
              r.max_trade_date=COALESCE(x.max_trade_date, r.max_trade_date),
              r.error_message=%s
            WHERE r.run_id=%s
            """,
            (run_id, status, error, run_id),
        )
    conn.commit()


def load_file(conn, path: Path, root: Path, archive: str, run_id: str) -> int:
    symbol, market = symbol_for(path)
    rel = str(path.relative_to(root)).replace("\\", "/")
    started = datetime.now()
    t0 = time.perf_counter()
    with conn.cursor() as cur:
        cur.execute("DROP TEMPORARY TABLE IF EXISTS tmp_full_adj_minute_csv")
        cur.execute(
            """
            CREATE TEMPORARY TABLE tmp_full_adj_minute_csv (
              timetag varchar(19) NOT NULL,
              open_price decimal(14,4) NULL,
              high_price decimal(14,4) NULL,
              low_price decimal(14,4) NULL,
              close_price decimal(14,4) NULL,
              volume decimal(24,4) NULL,
              amount decimal(24,4) NULL
            ) ENGINE=InnoDB
            """
        )
        cur.execute(
            """
            LOAD DATA LOCAL INFILE %s
            INTO TABLE tmp_full_adj_minute_csv
            CHARACTER SET utf8mb4
            FIELDS TERMINATED BY ',' OPTIONALLY ENCLOSED BY '"'
            LINES TERMINATED BY '\n'
            IGNORE 1 LINES
            (timetag, open_price, high_price, low_price, close_price, volume, amount)
            """,
            (str(path),),
        )
        loaded_tmp = cur.rowcount
        cur.execute(
            """
            INSERT INTO a_share_minute_price_full_adj
              (trade_date, trade_time, bar_time, symbol, market,
               open_price, high_price, low_price, close_price, volume, amount,
               source_file, source_archive, import_run_id)
            SELECT
              DATE(STR_TO_DATE(timetag, '%%Y%%m%%d %%H:%%i:%%s')),
              TIME(STR_TO_DATE(timetag, '%%Y%%m%%d %%H:%%i:%%s')),
              STR_TO_DATE(timetag, '%%Y%%m%%d %%H:%%i:%%s'),
              %s,
              %s,
              open_price,
              high_price,
              low_price,
              close_price,
              volume,
              amount,
              %s,
              %s,
              %s
            FROM tmp_full_adj_minute_csv
            ON DUPLICATE KEY UPDATE
              open_price=VALUES(open_price),
              high_price=VALUES(high_price),
              low_price=VALUES(low_price),
              close_price=VALUES(close_price),
              volume=VALUES(volume),
              amount=VALUES(amount),
              source_file=VALUES(source_file),
              source_archive=VALUES(source_archive),
              import_run_id=VALUES(import_run_id),
              imported_at=CURRENT_TIMESTAMP
            """,
            (symbol, market, rel, archive, run_id),
        )
        inserted = cur.rowcount
        duration_ms = int((time.perf_counter() - t0) * 1000)
        cur.execute(
            """
            INSERT INTO a_share_minute_full_adj_import_file_log
              (run_id, symbol, market, source_file, file_size, rows_loaded,
               started_at, finished_at, duration_ms, status)
            VALUES (%s, %s, %s, %s, %s, %s, %s, NOW(), %s, 'loaded')
            ON DUPLICATE KEY UPDATE
              rows_loaded=VALUES(rows_loaded),
              file_size=VALUES(file_size),
              finished_at=VALUES(finished_at),
              duration_ms=VALUES(duration_ms),
              status='loaded',
              error_message=NULL
            """,
            (run_id, symbol, market, rel, path.stat().st_size, loaded_tmp, started, duration_ms),
        )
        cur.execute(
            """
            UPDATE a_share_minute_full_adj_import_run
            SET files_loaded = (
              SELECT COUNT(*) FROM a_share_minute_full_adj_import_file_log
              WHERE run_id=%s AND status='loaded'
            )
            WHERE run_id=%s
            """,
            (run_id, run_id),
        )
    conn.commit()
    return loaded_tmp if loaded_tmp >= 0 else inserted


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--source-root", required=True)
    parser.add_argument("--source-archive", required=True)
    parser.add_argument("--run-id")
    parser.add_argument("--host", default="127.0.0.1")
    parser.add_argument("--port", type=int, default=3306)
    parser.add_argument("--user", default="root")
    parser.add_argument("--password")
    parser.add_argument("--database", default="tianxia")
    parser.add_argument("--limit", type=int)
    parser.add_argument("--skip-loaded", action="store_true")
    args = parser.parse_args()

    root = Path(args.source_root).resolve()
    files = discover_files(root)
    if args.limit:
        files = files[: args.limit]
    run_id = args.run_id or "full_adj_minute_" + datetime.now().strftime("%Y%m%d_%H%M%S")

    conn = connect(args)
    try:
        ensure_schema(conn)
        init_run(conn, run_id, args.source_archive, str(root), len(files))
        loaded_total = 0
        with conn.cursor() as cur:
            loaded_files = set()
            if args.skip_loaded:
                cur.execute(
                    "SELECT source_file FROM a_share_minute_full_adj_import_file_log WHERE run_id=%s AND status='loaded'",
                    (run_id,),
                )
                loaded_files = {row[0] for row in cur.fetchall()}
        for idx, path in enumerate(files, 1):
            rel = str(path.relative_to(root)).replace("\\", "/")
            if rel in loaded_files:
                continue
            rows = load_file(conn, path, root, args.source_archive, run_id)
            loaded_total += rows
            if idx % 10 == 0 or idx == 1 or idx == len(files):
                print(f"{datetime.now():%Y-%m-%d %H:%M:%S} loaded {idx}/{len(files)} files; last={rel}; rows={rows}; total_rows_seen={loaded_total}", flush=True)
        finish_run(conn, run_id, "completed")
        print(f"completed run_id={run_id} files={len(files)} rows_seen={loaded_total}")
        return 0
    except Exception as exc:
        conn.rollback()
        try:
            finish_run(conn, run_id, "failed", str(exc))
        except Exception:
            pass
        print(f"failed run_id={run_id}: {exc}", file=sys.stderr)
        return 1
    finally:
        conn.close()


if __name__ == "__main__":
    raise SystemExit(main())