from __future__ import annotations

import hashlib
import shutil
import sqlite3
import subprocess
from datetime import datetime, timedelta
from pathlib import Path
from sqlalchemy.orm import Session

from .config import settings
from .models import BackupRecord


def file_sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def postgres_cli_url(url: str) -> str:
    """PostgreSQL CLI 不识别 SQLAlchemy 的 +psycopg 驱动标记。"""
    return url.replace("postgresql+psycopg://", "postgresql://", 1)


def create_backup(db: Session, user_id: int | None = None) -> BackupRecord:
    stamp = datetime.now().astimezone().strftime("%Y%m%d-%H%M%S")
    url = settings.database_url
    if url.startswith("sqlite"):
        filename = f"database-{stamp}.sqlite3"
        target = settings.backup_dir / filename
        source_path = Path(url.split("///", 1)[1])
        if not source_path.is_absolute():
            source_path = Path.cwd() / source_path
        source = sqlite3.connect(str(source_path))
        destination = sqlite3.connect(str(target))
        try:
            source.backup(destination)
        finally:
            destination.close()
            source.close()
        database_type = "sqlite"
    elif url.startswith(("postgresql", "postgres")):
        filename = f"database-{stamp}.dump"
        target = settings.backup_dir / filename
        result = subprocess.run(["pg_dump", "--format=custom", "--file", str(target), postgres_cli_url(url)], capture_output=True, text=True)
        if result.returncode:
            raise RuntimeError(f"pg_dump 失败：{result.stderr.strip()}")
        database_type = "postgresql"
    else:
        raise RuntimeError("当前数据库类型暂不支持自动备份")
    record = BackupRecord(filename=filename, database_type=database_type, size_bytes=target.stat().st_size, sha256=file_sha256(target), created_by=user_id)
    db.add(record)
    db.commit()
    db.refresh(record)
    cleanup_old_backups(db)
    return record


def cleanup_old_backups(db: Session) -> None:
    cutoff = datetime.now().astimezone() - timedelta(days=settings.backup_keep_days)
    for record in db.query(BackupRecord).filter(BackupRecord.created_at < cutoff).all():
        path = settings.backup_dir / record.filename
        if path.exists():
            path.unlink()
        db.delete(record)
    db.commit()


def restore_backup(path: Path) -> None:
    if not path.resolve().is_relative_to(settings.backup_dir.resolve()):
        raise ValueError("备份文件必须位于 backups 目录")
    url = settings.database_url
    if url.startswith("sqlite"):
        source = sqlite3.connect(str(path))
        raw_path = Path(url.split("///", 1)[1])
        target_path = raw_path if raw_path.is_absolute() else Path.cwd() / raw_path
        temporary = target_path.with_suffix(".restore.tmp")
        destination = sqlite3.connect(str(temporary))
        try:
            source.backup(destination)
        finally:
            destination.close()
            source.close()
        shutil.move(temporary, target_path)
    elif url.startswith(("postgresql", "postgres")):
        result = subprocess.run(["pg_restore", "--clean", "--if-exists", "--no-owner", "--dbname", postgres_cli_url(url), str(path)], capture_output=True, text=True)
        if result.returncode:
            raise RuntimeError(f"pg_restore 失败：{result.stderr.strip()}")
    else:
        raise RuntimeError("当前数据库类型暂不支持恢复")