import sqlite3
import pandas as pd
from datetime import datetime

DB_NAME = "cognidetect.db"

def init_db():
    conn = sqlite3.connect(DB_NAME)
    cursor = conn.cursor()
    cursor.execute("""
        CREATE TABLE IF NOT EXISTS user_sessions (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            timestamp TEXT,
            age INTEGER, gender TEXT, country TEXT, grade TEXT,
            questionnaire_answers TEXT,
            questionnaire_risk_score TEXT,
            nlp_q_and_a TEXT,
            nlp_diagnosis TEXT,
            final_diagnosis TEXT
        )
    """)
    conn.commit()
    conn.close()

def save_session(user_data, rf_answers, rf_risk, nlp_data=None, final_diag="N/A"):
    conn = sqlite3.connect(DB_NAME)
    cursor = conn.cursor()
    
    # Format NLP data as string or N/A
    nlp_str = "N/A"
    nlp_diag = "N/A"
    if nlp_data:
        nlp_str = " | ".join([f"Q: {item['q']} A: {item['a']}" for item in nlp_data])
        nlp_diag = nlp_data[-1].get('diagnosis', 'N/A') if nlp_data else "N/A"

    cursor.execute("""
        INSERT INTO user_sessions 
        (timestamp, age, gender, country, grade, questionnaire_answers, questionnaire_risk_score, nlp_q_and_a, nlp_diagnosis, final_diagnosis)
        VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
    """, (
        datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
        user_data.get('age'), user_data.get('gender'), user_data.get('country'), user_data.get('grade'),
        str(rf_answers), str(rf_risk), nlp_str, nlp_diag, final_diag
    ))
    conn.commit()
    conn.close()

def get_all_data_as_csv():
    conn = sqlite3.connect(DB_NAME)
    df = pd.read_sql_query("SELECT * FROM user_sessions", conn)
    conn.close()
    return df.to_csv(index=False).encode('utf-8')