import asyncio
from io import BytesIO
from threading import Lock

import gradio as gr
from fastapi import Body, Depends, FastAPI, HTTPException
from PIL import Image
from sqlalchemy import Column, Float, ForeignKey, Integer, String, create_engine
from sqlalchemy.orm import declarative_base, relationship, sessionmaker
from transformers import AutoModelForCausalLM, AutoTokenizer

# Database setup
DATABASE_URL = "sqlite:///./sin_city_rp.db"
engine = create_engine(DATABASE_URL)
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
Base = declarative_base()
db_lock = Lock()

# Define SQLAlchemy models
class Player(Base):
    __tablename__ = "players"
    id = Column(Integer, primary_key=True, index=True)
    username = Column(String, unique=True, index=True)
    email = Column(String, unique=True, index=True)
    password = Column(String)
    characters = relationship("Character", back_populates="player")

class Character(Base):
    __tablename__ = "characters"
    id = Column(Integer, primary_key=True, index=True)
    name = Column(String, index=True)
    player_id = Column(Integer, ForeignKey("players.id"))
    player = relationship("Player", back_populates="characters")
    level = Column(Integer, default=1)
    experience = Column(Integer, default=0)
    health = Column(Integer, default=100)
    items = relationship("Item", back_populates="character")

class Quest(Base):
    __tablename__ = "quests"
    id = Column(Integer, primary_key=True, index=True)
    name = Column(String, index=True)
    description = Column(String)
    reward = Column(Integer)

class Item(Base):
    __tablename__ = "items"
    id = Column(Integer, primary_key=True, index=True)
    name = Column(String, index=True)
    description = Column(String)
    value = Column(Float)
    character_id = Column(Integer, ForeignKey("characters.id"))
    character = relationship("Character", back_populates="items")

# Create tables
Base.metadata.create_all(engine)

# FastAPI setup
app = FastAPI()

# Dependency for getting the current player
def get_current_player(username: str = Body(...), password: str = Body(...)):
    with db_lock:
        session = SessionLocal()
        player = session.query(Player).filter(Player.username == username, Player.password == password).first()
        session.close()
        if not player:
            raise HTTPException(status_code=401, detail="Invalid username or password")
        return player

# Login endpoint
@app.post("/login")
async def login(player: Player = Depends(get_current_player)):
    return {"message": f"Welcome {player.username}!"}

# Load the DALL·E Mini model
model_name = "flax-community/dalle-mini"
model = AutoModelForCausalLM.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)

# Function to generate image from text prompt using DALL·E Mini
async def generate_image(prompt: str):
    inputs = tokenizer(prompt, return_tensors="pt")
    outputs = model.generate(**inputs)
    image = Image.fromarray(outputs[0].numpy())

    # Convert the image to a format that can be used in FastAPI/Gradio
    buf = BytesIO()
    image.save(buf, format='PNG')
    buf.seek(0)
    return buf

# Gradio Interface
def gradio_interface(prompt):
    response = asyncio.run(generate_image(prompt))
    return response

interface = gr.Interface(
    fn=gradio_interface,
    inputs=gr.Textbox(label="Enter prompt"),
    outputs=gr.Image(label="Generated Image")
)

if __name__ == "__main__":
    interface.launch()