import sqlite3
import json
import os
from fastapi import FastAPI, Body
from fastapi.middleware.cors import CORSMiddleware
from sentence_transformers import SentenceTransformer
import uvicorn

app = FastAPI()

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

BASE_DIR = os.path.dirname(os.path.abspath(__file__))
DB_FILE = os.path.join(BASE_DIR, "assets.db")
MODEL_DIR = os.path.join(BASE_DIR, "local_model_miniLM")

if not os.path.exists(MODEL_DIR):
    model = SentenceTransformer('all-MiniLM-L6-v2')
    model.save(MODEL_DIR)
else:
    model = SentenceTransformer(MODEL_DIR)

def init_db():
    conn = sqlite3.connect(DB_FILE)
    conn.execute("""
                 CREATE TABLE IF NOT EXISTS documents
                 (
                     id INTEGER PRIMARY KEY AUTOINCREMENT,
                     title TEXT UNIQUE,
                     content TEXT,
                     embedding TEXT
                 )
                 """)
    conn.commit()
    conn.close()

init_db()

@app.post("/save-document")
async def save_document(data: dict = Body(...)):
    title = data.get("title")
    content = data.get("content")

    text_to_vector = f"{title} {str(content)}"
    vector = model.encode(text_to_vector).tolist()

    conn = sqlite3.connect(DB_FILE)
    try:
        conn.execute("""
                     INSERT INTO documents (title, content, embedding)
                     VALUES (?, ?, ?) ON CONFLICT(title) DO
                     UPDATE SET
                         content=excluded.content,
                         embedding=excluded.embedding
                     """, (title, json.dumps(content), json.dumps(vector)))
        conn.commit()
        return {"status": "success"}
    except Exception as e:
        return {"status": "error", "message": str(e)}
    finally:
        conn.close()

@app.get("/load-document")
async def load_document(title: str = None):
    conn = sqlite3.connect(DB_FILE)
    if title:
        row = conn.execute("SELECT title, content FROM documents WHERE title = ?", (title,)).fetchone()
    else:
        row = conn.execute("SELECT title, content FROM documents ORDER BY id DESC LIMIT 1").fetchone()
    conn.close()

    if row:
        return {"title": row[0], "content": json.loads(row[1])}
    return {"error": "Nie znaleziono dokumentu"}

if __name__ == "__main__":
    uvicorn.run(app, host="127.0.0.1", port=8000)