import os
import sqlite3
from datetime import datetime

import gradio as gr
from groq import Groq

# =====================================
# GROQ SETUP
# =====================================

GROQ_API_KEY = os.environ.get("GROQ_API_KEY")

if not GROQ_API_KEY:
    raise ValueError("GROQ_API_KEY environment variable is not set.")

client = Groq(api_key=GROQ_API_KEY)

# =====================================
# DATABASE SETUP
# =====================================

def init_db():

    conn = sqlite3.connect("memory.db")
    cursor = conn.cursor()

    # Conversation archive
    cursor.execute("""
    CREATE TABLE IF NOT EXISTS conversations (
        id INTEGER PRIMARY KEY AUTOINCREMENT,
        timestamp TEXT,
        role TEXT,
        content TEXT
    )
    """)

    # Long-term facts
    cursor.execute("""
    CREATE TABLE IF NOT EXISTS facts (
        id INTEGER PRIMARY KEY AUTOINCREMENT,
        timestamp TEXT,
        fact TEXT UNIQUE
    )
    """)

    conn.commit()
    conn.close()


init_db()

# =====================================
# MEMORY FUNCTIONS
# =====================================

def save_message(role, content):

    conn = sqlite3.connect("memory.db")
    cursor = conn.cursor()

    cursor.execute(
        """
        INSERT INTO conversations
        (timestamp, role, content)
        VALUES (?, ?, ?)
        """,
        (
            datetime.now().isoformat(),
            role,
            content
        )
    )

    conn.commit()
    conn.close()


def save_fact(fact):

    conn = sqlite3.connect("memory.db")
    cursor = conn.cursor()

    cursor.execute(
        """
        INSERT OR IGNORE INTO facts
        (timestamp, fact)
        VALUES (?, ?)
        """,
        (
            datetime.now().isoformat(),
            fact
        )
    )

    conn.commit()
    conn.close()


def get_memory_context(limit=20):

    conn = sqlite3.connect("memory.db")
    cursor = conn.cursor()

    cursor.execute("""
        SELECT role, content
        FROM conversations
        ORDER BY id DESC
        LIMIT ?
    """, (limit,))

    rows = cursor.fetchall()

    conn.close()

    rows.reverse()

    memory_text = ""

    for role, content in rows:
        memory_text += f"{role}: {content}\n"

    return memory_text


def get_fact_context():

    conn = sqlite3.connect("memory.db")
    cursor = conn.cursor()

    cursor.execute("""
        SELECT fact
        FROM facts
        ORDER BY id
    """)

    rows = cursor.fetchall()

    conn.close()

    fact_text = ""

    for row in rows:
        fact_text += "- " + row[0] + "\n"

    return fact_text


# =====================================
# PERSONALITY
# =====================================

with open("memory/core_personality.txt", "r", encoding="utf-8") as f:
    CORE_PERSONALITY = f.read()

SYSTEM_PROMPT = f"""
You are Teeg, a digital clone modeled after the real person.

Your behavior, tone, preferences, personality, and reasoning should follow the personality description below.

--- PERSONALITY START ---
{CORE_PERSONALITY}
--- PERSONALITY END ---

Stay consistent with Teeg's style at all times.
"""


# =====================================
# FACT EXTRACTION
# =====================================

def extract_facts(user_message, assistant_reply):

    try:

        extraction_prompt = f"""
You are a memory extraction system.

Your job is to identify long-term facts worth remembering.

Remember things that are likely to matter weeks or months from now.

Remember:
- preferences
- projects
- goals
- recurring interests
- personal background
- values

Ignore:
- greetings
- temporary questions
- jokes
- one-off requests

Conversation:

User: {user_message}

Assistant: {assistant_reply}

Return one fact per line.

If nothing should be remembered, return:

NONE
"""

        response = client.chat.completions.create(
            model="llama-3.3-70b-versatile",
            messages=[
                {
                    "role": "system",
                    "content": extraction_prompt
                }
            ],
            temperature=0
        )

        result = response.choices[0].message.content.strip()

        if result.upper() == "NONE":
            return

        facts = result.split("\n")

        for fact in facts:

            fact = fact.strip()

            if fact:
                save_fact(fact)

    except Exception as e:
        print("Fact extraction error:", e)


# =====================================
# CHAT
# =====================================

def chat(message, history):

    memory_context = get_memory_context()
    fact_context = get_fact_context()

    messages = [
        {
            "role": "system",
            "content":
                SYSTEM_PROMPT
                + "\n\nKNOWN FACTS:\n"
                + fact_context
                + "\nRECENT MEMORIES:\n"
                + memory_context
        }
    ]

    # Current browser session
    for msg in history:

        role = msg.get("role")
        content = msg.get("content")

        if role in ["user", "assistant"] and content:

            messages.append({
                "role": role,
                "content": content
            })

    messages.append({
        "role": "user",
        "content": message
    })

    save_message(
        "user",
        message
    )

    try:

        response = client.chat.completions.create(
            model="llama-3.3-70b-versatile",
            messages=messages,
            temperature=0.7
        )

        reply = response.choices[0].message.content

        save_message(
            "assistant",
            reply
        )

        # Learn from conversation
        extract_facts(
            message,
            reply
        )

        return reply

    except Exception as e:
        return f"Error: {str(e)}"


# =====================================
# UI
# =====================================

demo = gr.ChatInterface(
    fn=chat,
    title="🧬 Teeg Clone",
    description="A personality-driven AI clone powered by Groq"
)

demo.launch()