import os
import asyncio
import logging
from dotenv import load_dotenv
from nio import AsyncClient, AsyncClientConfig, MatrixRoom, RoomMessageText, InviteMemberEvent
from nio.responses import LoginResponse
from openai import AsyncOpenAI

# --- Load environment variables ---
load_dotenv()
HOMESERVER_URL = os.getenv("HOMESERVER_URL")
USER_ID = os.getenv("USER_ID")
PASSWORD = os.getenv("PASSWORD")
LOG_LEVEL = os.getenv("LOG_LEVEL", "INFO").upper()
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")

if not OPENAI_API_KEY:
    raise RuntimeError("OPENAI_API_KEY is not set in environment")

# --- Initialize Async OpenAI client ---
openai_client = AsyncOpenAI(api_key=OPENAI_API_KEY)

# --- Logging Setup ---
numeric_level = getattr(logging, LOG_LEVEL, logging.INFO)
logging.basicConfig(
    level=numeric_level,
    format="%(asctime)s %(levelname)s %(name)s: %(message)s"
)
logger = logging.getLogger(__name__)

async def trust_all_devices(client) -> None:
    """
    Programmatically verify all devices to allow sharing encryption keys.
    """
    for room_id in client.rooms:
        try:
            devices = await client.room_devices(room_id)
            if isinstance(devices, dict):
                for user, dev_ids in devices.items():
                    if user == USER_ID:
                        continue
                    for dev_id in dev_ids:
                        device = client.crypto.device_store.get_device(user, dev_id)
                        if device and not client.crypto.device_store.is_device_verified(device):
                            logger.info(f"Trusting {dev_id} for {user}")
                            client.verify_device(device)
        except Exception:
            logger.exception(f"Error trusting devices in {room_id}")

async def message_callback(room: MatrixRoom, event: RoomMessageText):
    """Handle incoming text messages."""
    if event.sender == USER_ID:
        return
    body = event.body.strip()
    lower = body.lower()
    logger.info("Received '%s' from %s in %s", body, event.sender, room.display_name)

    send_kwargs = {
        "room_id": room.room_id,
        "message_type": "m.room.message",
        "ignore_unverified_devices": True
    }

    # Simple ping
    if lower == "!ping":
        await client.room_send(**send_kwargs, content={"msgtype": "m.text", "body": "Pong!"})
        return

    # Ask OpenAI via chat completion
    if lower.startswith("!ask "):
        question = body[5:].strip()
        if not question:
            await client.room_send(**send_kwargs, content={"msgtype": "m.text", "body": "Provide a question after !ask."})
            return
        logger.info("Querying OpenAI: %s", question)
        try:
            response = await openai_client.chat.completions.create(
                model="gpt-3.5-turbo",
                messages=[
                    {"role": "system", "content": "You are a helpful assistant."},
                    {"role": "user",   "content": question}
                ],
                max_tokens=150
            )
            answer = response.choices[0].message.content.strip()
        except Exception:
            logger.exception("OpenAI API error")
            answer = "Sorry, I encountered an error contacting the AI service."
        await client.room_send(**send_kwargs, content={"msgtype": "m.text", "body": answer})
        return

    # Greeting
    if lower == "hello botbot":
        await client.room_send(**send_kwargs, content={"msgtype": "m.text", "body": "Hello! How can I assist you today?"})

async def main() -> None:
    """Initialize and run the Matrix bot."""
    global client
    config = AsyncClientConfig(store_sync_tokens=True, encryption_enabled=True)
    client = AsyncClient(HOMESERVER_URL, USER_ID, store_path="/app/data", config=config)

    login_resp = await client.login(password=PASSWORD)
    if isinstance(login_resp, LoginResponse):
        logger.info("Logged in as %s", USER_ID)
    else:
        logger.error("Login failed: %s", login_resp)
        return

    await trust_all_devices(client)

    # Auto-join and trust
    async def on_invite(room, event):
        if isinstance(event, InviteMemberEvent):
            await client.join(room.room_id)
            logger.info("Joined %s", room.room_id)
            await trust_all_devices(client)
    client.add_event_callback(on_invite, InviteMemberEvent)
    client.add_event_callback(message_callback, RoomMessageText)

    logger.info("Starting sync loop")
    await client.sync_forever(timeout=30000)

if __name__ == "__main__":
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        logger.info("Shutting down")
        asyncio.run(client.close())