import os
from contextlib import closing
import pymysql

from .mysql_client import get_mysql_connection

TABLE = "instance_usage"


def _connect():
    # Reuse central connection settings (AWS/local env) from mysql_client
    return get_mysql_connection()


def ensure_instance_row(instance_id: str) -> bool:
    """Ensure an instance_usage row exists. Non-destructive: does not modify user/company.

    Returns True if the row exists (created or already present).
    """
    if not instance_id:
        return False
    sql = f"""
        INSERT INTO {TABLE} (instance_id, status)
        VALUES (%s, 'available')
        ON DUPLICATE KEY UPDATE instance_id = instance_id
    """
    with closing(_connect()) as conn, conn.cursor() as cur:
        cur.execute(sql, (instance_id,))
        try:
            conn.commit()
        except Exception:
            pass
        return True


def acquire_lock(instance_id: str, user_id: str, company_id: str | None = None) -> bool:
    """Mark this instance as in use by this user. Idempotent if already held by same user."""
    # Ensure the instance row exists first (non-destructive)
    try:
        ensure_instance_row(instance_id)
    except Exception:
        pass
    sql = f"""
        UPDATE {TABLE}
        SET status='in_use', user_id=%s, company_id=%s, locked_at=UTC_TIMESTAMP(6)
        WHERE instance_id=%s AND (status='available' OR user_id=%s)
    """
    with closing(_connect()) as conn, conn.cursor() as cur:
        rows = cur.execute(sql, (user_id, company_id, instance_id, user_id))
        try:
            conn.commit()
        except Exception:
            pass
        return rows > 0


def send_heartbeat(instance_id: str, user_id: str) -> bool:
    """Refresh locked_at only if the lock is still held by this user."""
    sql = f"""
        UPDATE {TABLE}
        SET locked_at=UTC_TIMESTAMP(6)
        WHERE instance_id=%s AND user_id=%s AND status='in_use'
    """
    with closing(_connect()) as conn, conn.cursor() as cur:
        rows = cur.execute(sql, (instance_id, user_id))
        try:
            conn.commit()
        except Exception:
            pass
        return rows > 0


def release_lock(instance_id: str, user_id: str) -> None:
    """Release the lock when the job/session ends."""
    sql = f"""
        UPDATE {TABLE}
        SET status='available', user_id=NULL, company_id=NULL, locked_at=NULL
        WHERE instance_id=%s AND user_id=%s
    """
    with closing(_connect()) as conn, conn.cursor() as cur:
        cur.execute(sql, (instance_id, user_id))
        try:
            conn.commit()
        except Exception:
            pass