import uuid
from datetime import datetime, timezone
from pymongo import ReturnDocument
from pymongo.asynchronous.database import AsyncDatabase
from app.utils.country_resolver import resolve_country

async def create_session(
    db: AsyncDatabase,
    site_id: str,
    user_data: dict | None = None,
    ip_address: str | None = None
) -> dict:
    session_id = f"sess_{uuid.uuid4().hex[:12]}"
    now = datetime.now(timezone.utc)
    country = resolve_country(ip_address)
    
    session_doc = {
        "sessionId": session_id,
        "siteId": site_id,
        "mode": "rule",
        "status": "RINGING",
        "assignedAgentId": None,
        "assignedAt": None,
        "connectionStatus": "online",
        "disconnectedAt": None,
        "userId": user_data.get("id") if user_data else None,
        "userName": user_data.get("name") if user_data else None,
        "userEmail": user_data.get("email") if user_data else None,
        "ipAddress": ip_address,
        "country": country,
        "createdAt": now,
        "lastActiveAt": now
    }
    
    await db["sessions"].insert_one(session_doc)
    return session_doc

async def get_session_by_id(db: AsyncDatabase, session_id: str) -> dict | None:
    return await db["sessions"].find_one({"sessionId": session_id})

async def update_session_status(db: AsyncDatabase, session_id: str, status_val: str) -> None:
    now = datetime.now(timezone.utc)
    mode_val = "human" if status_val in ("IN_CHAT", "ENDED") else "rule"
    await db["sessions"].update_one(
        {"sessionId": session_id},
        {"$set": {"status": status_val, "mode": mode_val, "lastActiveAt": now}}
    )

async def update_session_connection_status(
    db: AsyncDatabase,
    session_id: str,
    connection_status: str,
    disconnected_at: datetime | None = None
) -> None:
    now = datetime.now(timezone.utc)
    await db["sessions"].update_one(
        {"sessionId": session_id},
        {"$set": {"connectionStatus": connection_status, "disconnectedAt": disconnected_at, "lastActiveAt": now}}
    )

async def claim_session_atomically(db: AsyncDatabase, session_id: str, agent_id: str) -> dict | None:
    """
    Atomically claims a session for an agent only if it is currently in RINGING state.
    Returns the updated session document if claim succeeded, or None if failed/already claimed.
    """
    if not session_id or not agent_id:
        return None
    now = datetime.now(timezone.utc)
    return await db["sessions"].find_one_and_update(
        {"sessionId": session_id, "status": "RINGING"},
        {
            "$set": {
                "status": "IN_CHAT",
                "mode": "human",
                "assignedAgentId": agent_id,
                "assignedAt": now,
                "lastActiveAt": now
            }
        },
        return_document=ReturnDocument.AFTER
    )

async def requeue_session(db: AsyncDatabase, session_id: str, agent_id: str | None = None) -> dict | None:
    """
    Safely requeues an IN_CHAT session back to RINGING state and clears agent assignment.
    If agent_id is provided, only requeues if currently assigned to that agent.
    Returns the updated session document if requeued, or None if invalid state/already claimed.
    """
    if not session_id:
        return None
    now = datetime.now(timezone.utc)
    query = {"sessionId": session_id, "status": "IN_CHAT"}
    if agent_id:
        query["assignedAgentId"] = agent_id

    return await db["sessions"].find_one_and_update(
        query,
        {
            "$set": {
                "status": "RINGING",
                "mode": "rule",
                "lastActiveAt": now
            },
            "$unset": {
                "assignedAgentId": "",
                "assignedAt": ""
            }
        },
        return_document=ReturnDocument.AFTER
    )