import os
import uuid
import datetime
import gradio as gr
from transformers import pipeline
import spaces
import torch
import csv

DATA_DIR = os.path.join(os.getcwd(), "session_data")
os.makedirs(DATA_DIR, exist_ok=True)

pipe = pipeline("text-generation", model="mecoffey/NPC_brain", trust_remote_code=True)

SYSTEM_PROMPT = (
    "You are a creative storyteller. "
    "Write vivid, engaging short stories with clear structure and imaginative detail."
    " Keep Stories to a maximum 5 paragraphs."
)


def make_session_csv():
    session_id = uuid.uuid4().hex
    return os.path.join(DATA_DIR, f"story_session_{session_id}.csv")


def append_to_csv(csv_path, prompt, story):
    is_new = not os.path.exists(csv_path)
    with open(csv_path, "a", newline="", encoding="utf-8") as f:
        writer = csv.writer(f)
        if is_new:
            writer.writerow(["timestamp", "prompt", "story"])
        writer.writerow([
            datetime.datetime.utcnow().isoformat(),
            prompt,
            story,
        ])


def read_csv_rows(csv_path):
    if not csv_path or not os.path.exists(csv_path):
        return []
    with open(csv_path, newline="", encoding="utf-8") as f:
        reader = csv.reader(f)
        rows = list(reader)
    return rows[1:] if len(rows) > 1 else []


@spaces.GPU
def ask(message, session_csv):
    if session_csv is None:
        session_csv = make_session_csv()

    prompt = f"{message}\n\nDescribe:"
    response = pipe(prompt, return_full_text=False, max_new_tokens=256)
    story = response[0]["generated_text"].strip()

    append_to_csv(session_csv, message, story)
    csv_rows = read_csv_rows(session_csv)

    return story, csv_rows, session_csv


with gr.Blocks() as demo:
    text_in = gr.Textbox(label="Story Idea", placeholder="A girl in a red hood...")
    text_out = gr.Textbox(label="Story", interactive=False)
    csv_table = gr.Dataframe(
        headers=["Timestamp", "Prompt", "Story"],
        interactive=False,
        row_count=(1, "dynamic"),
    )
    session_csv_state = gr.State(value=None)
    btn = gr.Button("Write!")

    btn.click(
        fn=ask,
        inputs=[text_in, session_csv_state],
        outputs=[text_out, csv_table, session_csv_state],
    )


demo.launch()