from sqlmodel import SQLModel, Field, create_engine, Session, select
from datetime import datetime
from typing import Optional
import os
import json
from pptx import Presentation
import fitz  # PyMuPDF
from openai import OpenAI

# === Setup ===
db_path = "/tmp/chatter_sessions.db"
engine = create_engine(f"sqlite:///{db_path}")
SQLModel.metadata.create_all(engine)

openai_api_key = os.getenv("OPENAI_API_KEY")
client = OpenAI(api_key=openai_api_key)

# === DB Tables ===
class StudyGuideEntry(SQLModel, table=True):
    id: Optional[int] = Field(default=None, primary_key=True)
    user: str
    filename: str
    guide: str
    timestamp: str

class UserFileEntry(SQLModel, table=True):
    id: Optional[int] = Field(default=None, primary_key=True)
    user: str
    filename: str
    file_path: str
    timestamp: str

class UserProfile(SQLModel, table=True):
    id: Optional[int] = Field(default=None, primary_key=True)
    nickname: str
    timestamp: str

SQLModel.metadata.create_all(engine)

# === File Persistence ===
def save_user_file(user, filename, file_path):
    session = Session(engine)
    entry = UserFileEntry(
        user=user,
        filename=filename,
        file_path=file_path,
        timestamp=datetime.now().strftime("%Y-%m-%d %H:%M")
    )
    session.add(entry)
    session.commit()
    session.close()

def fetch_user_files(user):
    session = Session(engine)
    statement = select(UserFileEntry).where(UserFileEntry.user == user)
    results = session.exec(statement).all()
    session.close()
    return results

def save_study_guide(user, filename, guide):
    session = Session(engine)
    entry = StudyGuideEntry(
        user=user,
        filename=filename,
        guide=guide,
        timestamp=datetime.now().strftime("%Y-%m-%d %H:%M")
    )
    session.add(entry)
    session.commit()
    session.close()

def fetch_study_guides(user):
    session = Session(engine)
    statement = select(StudyGuideEntry).where(StudyGuideEntry.user == user)
    results = session.exec(statement).all()
    session.close()
    return results

def save_nickname_to_db(nickname):
    session = Session(engine)
    existing = session.exec(select(UserProfile).where(UserProfile.nickname == nickname)).first()
    if not existing:
        entry = UserProfile(nickname=nickname, timestamp=datetime.now().strftime("%Y-%m-%d %H:%M"))
        session.add(entry)
        session.commit()
    session.close()

def load_latest_nickname():
    session = Session(engine)
    result = session.exec(select(UserProfile).order_by(UserProfile.id.desc())).first()
    session.close()
    return result.nickname if result else ""

# === LLM Core ===
def call_llm(prompt=None, system_message="You are a helpful assistant.", messages=None):
    max_safe_chars = 28000
    if messages:
        for msg in messages:
            if len(msg["content"]) > max_safe_chars:
                msg["content"] = msg["content"][:max_safe_chars] + "\n\n[Truncated due to token limit]"
    else:
        if prompt and len(prompt) > max_safe_chars:
            prompt = prompt[:max_safe_chars] + "\n\n[Truncated due to token limit]"
        messages = [
            {"role": "system", "content": system_message},
            {"role": "user", "content": prompt}
        ]

    response = client.chat.completions.create(
        model="gpt-4",
        messages=messages,
        temperature=0.7
    )
    return response.choices[0].message.content

# === File Parsing ===
def extract_text_from_file(file):
    ext = os.path.splitext(file.name)[1].lower()
    if ext == ".pdf":
        with fitz.open(file.name) as doc:
            return "\n".join([page.get_text() for page in doc])
    elif ext in [".txt", ".md"]:
        return file.read().decode("utf-8")
    elif ext == ".pptx":
        prs = Presentation(file.name)
        return "\n".join([
            shape.text for slide in prs.slides for shape in slide.shapes if hasattr(shape, "text")
        ])
    else:
        return "Unsupported file type."

# === GPT-Based Generators ===
def generate_summary(text):
    prompt = f"Summarize the following document in 5–7 bullet points:\n\n{text}"
    return call_llm(prompt=prompt)

def generate_flashcards(text):
    prompt = (
        "Generate 5 flashcards based on the document below.\n"
        "Each flashcard should follow this format:\nQ: <question>\nA: <answer>\n\n"
        f"{text}"
    )
    return call_llm(prompt=prompt)

def generate_quiz(text):
    prompt = (
        "Generate 5 multiple choice questions from the document below. "
        "Return a JSON array where each item has:\n"
        "- question (string)\n- options (list of strings)\n- answer (correct option string)\n\n"
        f"{text}"
    )
    response = call_llm(prompt=prompt)
    try:
        return json.loads(response)
    except:
        return []

def answer_question(text, question):
    prompt = f"Using only the document below, answer this question:\n\nDocument:\n{text}\n\nQuestion:\n{question}"
    return call_llm(prompt=prompt)