#!/usr/bin/env python3

import os
import sys
import mysql.connector
from mysql.connector import errorcode

CFG = {
    "host": os.getenv("DB_HOST", "127.0.0.1"),
    "user": os.getenv("DB_USER", "cs348"),
    "password": os.getenv("DB_PASS", "Cs348Group!"),
    "database": os.getenv("DB_NAME", "cs348db"),
    "charset": "utf8mb4",
    "collation": "utf8mb4_unicode_ci",
    "auth_plugin": "caching_sha2_password",
}

DDL = """
CREATE TABLE IF NOT EXISTS `ChatHistory` (
  `text_id` INT NOT NULL,
  `type`    VARCHAR(64) NOT NULL,
  `text`    MEDIUMTEXT NOT NULL,
  PRIMARY KEY (`text_id`)
) ENGINE=InnoDB
  DEFAULT CHARSET = utf8mb4
  COLLATE = utf8mb4_unicode_ci;
"""

def get_connection():
    try:
        return mysql.connector.connect(**CFG)
    except mysql.connector.Error as e:
        print(f"[fatal] MySQL connect error: {e}", file=sys.stderr)
        sys.exit(1)

def ensure_table(conn):
    with conn.cursor() as cur:
        cur.execute(DDL)
    conn.commit()

def get_next_id(conn):
    with conn.cursor() as cur:
        cur.execute("SELECT COALESCE(MAX(text_id), 0) FROM `ChatHistory`;")
        (max_id,) = cur.fetchone()
    return int(max_id) + 1

def insert_two_rows(conn, user_text: str):
    nxt = get_next_id(conn)
    rows = [
        (nxt,     "user",      user_text),
        (nxt + 1, "assistant", "LLM answer here"),
    ]
    with conn.cursor() as cur:
        cur.executemany(
            "INSERT INTO `ChatHistory` (`text_id`, `type`, `text`) VALUES (%s, %s, %s);",
            rows,
        )
    conn.commit()
    return nxt, nxt + 1

def list_history(conn, limit: int | None = None):
    sql = "SELECT `text_id`, `type`, `text` FROM `ChatHistory` ORDER BY `text_id` ASC"
    params = ()
    if limit is not None and limit > 0:
        sql += " LIMIT %s"
        params = (limit,)
    with conn.cursor() as cur:
        cur.execute(sql, params)
        rows = cur.fetchall()
    if not rows:
        print("(no rows)")
        return
    # Simple pretty print without extra deps
    idw = max(len("text_id"), *(len(str(r[0])) for r in rows))
    typew = max(len("type"), *(len(str(r[1])) for r in rows))
    print(f"{'text_id'.ljust(idw)}  {'type'.ljust(typew)}  text")
    print(f"{'-'*idw}  {'-'*typew}  {'-'*4}")
    for tid, typ, txt in rows:
        print(f"{str(tid).ljust(idw)}  {str(typ).ljust(typew)}  {txt}")

def print_help():
    print(r"""
Commands:
  \list            Print all history (ordered by text_id)
  \list N          Print first N rows (ordered by text_id)
  \q               Quit
  \help            This help

Anything else is sent as a user prompt:
  - Inserts (next_id,   'user',      <your text>)
  - Inserts (next_id+1, 'assistant', 'LLM answer here')
""".strip())

def main():
    print("Connecting to MySQL…")
    conn = get_connection()
    try:
        ensure_table(conn)
        print("Ready. Type \\help for commands. Empty line quits.")
        while True:
            try:
                line = input("Enter prompt> ")
            except (EOFError, KeyboardInterrupt):
                print("\nBye.")
                break

            if line is None:
                continue
            line = line.strip()
            if not line:
                print("Bye.")
                break

            # Commands
            if line == r"\q":
                print("Bye.")
                break
            if line == r"\help":
                print_help()
                continue
            if line.startswith(r"\list"):
                parts = line.split()
                limit = None
                if len(parts) == 2 and parts[1].isdigit():
                    limit = int(parts[1])
                try:
                    list_history(conn, limit)
                except mysql.connector.Error as e:
                    print(f"[error] list failed: {e}", file=sys.stderr)
                continue

            # Regular insert path
            try:
                first_id, second_id = insert_two_rows(conn, line)
                print(f"Inserted prompt at text_id={first_id} and placeholder at text_id={second_id}.")
            except mysql.connector.Error as e:
                print(f"[error] insert failed: {e}", file=sys.stderr)
                if e.errno in (errorcode.CR_SERVER_GONE_ERROR, errorcode.CR_SERVER_LOST):
                    try:
                        conn.close()
                    except Exception:
                        pass
                    conn = get_connection()
    finally:
        try:
            conn.close()
        except Exception:
            pass

if __name__ == "__main__":
    main()