import os
import sqlite3
from datetime import datetime
import gradio as gr
from huggingface_hub import InferenceClient
from datasets import load_dataset

# ---------------------------
# Config
# ---------------------------
MODELS = {
    "Meta LLaMA 3.1 (8B Instruct)": "meta-llama/Llama-3.1-8B-Instruct",
    "Mistral 7B Instruct": "mistralai/Mistral-7B-Instruct-v0.3",
}

DATASETS = ["The Stack", "CodeXGLUE"]
HF_TOKEN = os.getenv("HF_TOKEN")
DB_PATH = "history.db"

SYSTEM_DEFAULT = (
    "You are a backend-focused coding assistant. "
    "Always prioritize database, API, authentication, routing, migrations, and CRUD logic. "
    "Provide full backend code scaffolds with files, paths, and commands. "
    "Only include frontend if required for framework integration."
)

# ---------------------------
# DB Setup (unchanged)
# ---------------------------
def db():
    conn = sqlite3.connect(DB_PATH)
    conn.execute("PRAGMA journal_mode=WAL;")
    return conn

def init_db():
    conn = db()
    cur = conn.cursor()
    cur.execute("""
        CREATE TABLE IF NOT EXISTS sessions (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            title TEXT NOT NULL,
            created_at TEXT NOT NULL
        )
    """)
    cur.execute("""
        CREATE TABLE IF NOT EXISTS messages (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            session_id INTEGER NOT NULL,
            role TEXT NOT NULL,
            content TEXT NOT NULL,
            created_at TEXT NOT NULL,
            FOREIGN KEY(session_id) REFERENCES sessions(id) ON DELETE CASCADE
        )
    """)
    conn.commit()
    conn.close()

def create_session(title: str = "New chat") -> int:
    conn = db()
    cur = conn.cursor()
    cur.execute(
        "INSERT INTO sessions (title, created_at) VALUES (?, ?)",
        (title, datetime.utcnow().isoformat())
    )
    session_id = cur.lastrowid
    conn.commit()
    conn.close()
    return session_id

def delete_session(session_id: int):
    conn = db()
    cur = conn.cursor()
    cur.execute("DELETE FROM messages WHERE session_id = ?", (session_id,))
    cur.execute("DELETE FROM sessions WHERE id = ?", (session_id,))
    conn.commit()
    conn.close()

def list_sessions():
    conn = db()
    cur = conn.cursor()
    cur.execute("SELECT id, title FROM sessions ORDER BY id DESC")
    rows = cur.fetchall()
    conn.close()
    labels = [f"{r[0]} • {r[1]}" for r in rows]
    return labels, rows

def get_messages(session_id: int):
    conn = db()
    cur = conn.cursor()
    cur.execute("""
        SELECT role, content
        FROM messages
        WHERE session_id = ?
        ORDER BY id ASC
    """, (session_id,))
    rows = cur.fetchall()
    conn.close()
    msgs = [{"role": role, "content": content} for (role, content) in rows]
    return msgs

def add_message(session_id: int, role: str, content: str):
    conn = db()
    cur = conn.cursor()
    cur.execute(
        "INSERT INTO messages (session_id, role, content, created_at) VALUES (?, ?, ?, ?)",
        (session_id, role, content, datetime.utcnow().isoformat())
    )
    conn.commit()
    conn.close()

def update_session_title_if_needed(session_id: int, first_user_text: str):
    conn = db()
    cur = conn.cursor()
    cur.execute("SELECT COUNT(*) FROM messages WHERE session_id=? AND role='user'", (session_id,))
    count_users = cur.fetchone()[0]
    if count_users == 1:
        title = first_user_text.strip().split("\n")[0]
        title = (title[:50] + "…") if len(title) > 50 else title
        cur.execute("UPDATE sessions SET title=? WHERE id=?", (title or "New chat", session_id))
        conn.commit()
    conn.close()

# ---------------------------
# Helpers (unchanged)
# ---------------------------
def label_to_id(label: str | None) -> int | None:
    if not label:
        return None
    try:
        return int(label.split("•", 1)[0].strip())
    except Exception:
        return None

def build_api_messages(session_id: int, system_message: str):
    msgs = [{"role": "system", "content": system_message.strip()}]
    msgs.extend(get_messages(session_id))
    return msgs

def get_client(model_choice: str):
    model_id = MODELS.get(model_choice, list(MODELS.values())[0])
    return InferenceClient(model_id, token=HF_TOKEN)

def load_dataset_by_name(name: str):
    if name == "The Stack":
        return load_dataset("bigcode/the-stack", split="train")
    elif name == "CodeXGLUE":
        return load_dataset("google/code_x_glue_cc_code_to_code_trans", split="train")
    return None

# ---------------------------
# Gradio Callbacks (unchanged)
# ---------------------------
def refresh_sessions_cb():
    labels, _ = list_sessions()
    selected = labels[0] if labels else None
    return gr.update(choices=labels, value=selected)

def new_chat_cb():
    sid = create_session("New chat")
    labels, _ = list_sessions()
    selected = next((lbl for lbl in labels if lbl.startswith(f"{sid} ")), None)
    return (gr.update(choices=labels, value=selected), [], "")

def load_session_cb(selected_label):
    sid = label_to_id(selected_label)
    if not sid:
        return []
    return get_messages(sid)

def delete_chat_cb(selected_label):
    sid = label_to_id(selected_label)
    if sid:
        delete_session(sid)
    labels, _ = list_sessions()
    selected = labels[0] if labels else None
    return gr.update(choices=labels, value=selected), []

FRONTEND_KEYWORDS = [
    "react", "vue", "angular", "html", "css", "javascript", "tailwind", "recharts", "typescript"
]

def is_frontend_request(user_text: str) -> bool:
    text_lower = user_text.lower()
    return any(kw in text_lower for kw in FRONTEND_KEYWORDS)

def send_cb(user_text, selected_label, chatbot_msgs, system_message, max_tokens, temperature, top_p, model_choice, dataset_choice, *args):
    sid = label_to_id(selected_label)
    if sid is None:
        sid = create_session("New chat")
        labels, _ = list_sessions()
        selected_label = next((lbl for lbl in labels if lbl.startswith(f"{sid} ")), None)

    add_message(sid, "user", user_text)
    update_session_title_if_needed(sid, user_text)

    display_msgs = chatbot_msgs[:]
    display_msgs.append({"role": "user", "content": user_text})

    if is_frontend_request(user_text):
        apology = "⚠️ I'm a backend-focused assistant and cannot provide frontend code."
        display_msgs.append({"role": "assistant", "content": apology})
        add_message(sid, "assistant", apology)
        yield (display_msgs, "", selected_label)
        return

    display_msgs.append({"role": "assistant", "content": "…"})
    yield (display_msgs, "", selected_label)

    client = get_client(model_choice)
    api_messages = build_api_messages(sid, system_message)
    partial = ""

    try:
        for chunk in client.chat_completion(
            messages=api_messages,
            max_tokens=int(max_tokens),
            temperature=float(temperature),
            top_p=float(top_p),
            stream=True,
        ):
            if not hasattr(chunk, "choices") or not chunk.choices:
                continue
            choice = chunk.choices[0]
            delta = ""
            if hasattr(choice, "delta") and choice.delta and getattr(choice.delta, "content", None) is not None:
                delta = choice.delta.content
            elif hasattr(choice, "message") and getattr(choice.message, "content", None) is not None:
                delta = choice.message.content

            if delta:
                partial += delta
                display_msgs[-1]["content"] = partial
                yield (display_msgs, "", selected_label)

        add_message(sid, "assistant", partial)

    except Exception as e:
        display_msgs[-1]["content"] = f"⚠️ Error: {str(e)}"
        yield (display_msgs, "", selected_label)

def regenerate_cb(selected_label, system_message, max_tokens, temperature, top_p, model_choice, dataset_choice):
    sid = label_to_id(selected_label)
    if sid is None:
        return [], ""

    msgs = get_messages(sid)
    if not msgs:
        return [], ""

    if msgs and msgs[-1]["role"] == "assistant":
        conn = db()
        cur = conn.cursor()
        cur.execute("""
            DELETE FROM messages
            WHERE id = (
                SELECT id FROM messages
                WHERE session_id=?
                ORDER BY id DESC LIMIT 1
            )
        """, (sid,))
        conn.commit()
        conn.close()
        msgs = get_messages(sid)

    dataset = load_dataset_by_name(dataset_choice)
    api_messages = [{"role": "system", "content": system_message.strip()}] + msgs
    display_msgs = msgs + [{"role": "assistant", "content": ""}]

    client = get_client(model_choice)
    partial = ""

    try:
        for chunk in client.chat_completion(
            messages=api_messages,
            max_tokens=int(max_tokens),
            temperature=float(temperature),
            top_p=float(top_p),
            stream=True,
        ):
            if not hasattr(chunk, "choices") or not chunk.choices:
                continue
            choice = chunk.choices[0]
            delta = ""
            if hasattr(choice, "delta") and choice.delta and getattr(choice.delta, "content", None) is not None:
                delta = choice.delta.content
            elif hasattr(choice, "message") and getattr(choice.message, "content", None) is not None:
                delta = choice.message.content

            if delta:
                partial += delta
                display_msgs[-1]["content"] = partial
                yield display_msgs

        add_message(sid, "assistant", partial)

    except Exception as e:
        display_msgs[-1]["content"] = f"⚠️ Error: {str(e)}"
        yield display_msgs

# ---------------------------
# App UI - OPTIMIZED LAYOUT
# ---------------------------
init_db()
labels, _ = list_sessions()
if not labels:
    first_sid = create_session("New chat")
    labels, _ = list_sessions()

default_selected = labels[0] if labels else None

with gr.Blocks(title="Backend-Focused LLaMA/Mistral CRUD Assistant", theme=gr.themes.Soft()) as demo:
    gr.HTML("""
    <style>
    .compact-row {
        margin-bottom: 0.5rem !important;
    }
    .compact-slider {
        margin-bottom: 0.3rem !important;
    }
    .small-textbox {
        margin-bottom: 0.5rem !important;
    }
    .section-header {
        margin-bottom: 0.5rem !important;
        font-weight: 600;
    }
    button {
        background-color: #22c55e !important;
        color: #ffffff !important;
        border: none !important;
    }
    button:hover {
        background-color: #16a34a !important;
    }
    button:focus {
        outline: 2px solid #166534 !important;
        outline-offset: 2px;
    }
    </style>
    """)

    gr.Markdown("## 🗄️ Backend CRUD Assistant")

    with gr.Row():
        # Left sidebar - compact layout
        with gr.Column(scale=1, min_width=280):
            # Sessions section
            gr.Markdown("### 📁 Sessions", elem_classes=["section-header"])
            session_list = gr.Radio(
                choices=labels,
                value=default_selected,
                label="Your chats",
                interactive=True,
                elem_classes=["compact-row"]
            )
            
            # Session management buttons in a tight grid
            with gr.Row(elem_classes=["compact-row"]):
                new_btn = gr.Button("➕ New", size="sm", variant="primary")
                del_btn = gr.Button("🗑️ Delete", size="sm", variant="stop")
                refresh_btn = gr.Button("🔄", size="sm", variant="secondary")
            
            # Rename section - more compact
            with gr.Row(elem_classes=["compact-row"]):
                edit_title_box = gr.Textbox(
                    label="✏️ Rename", 
                    placeholder="New title...",
                    scale=3,
                    elem_classes=["small-textbox"]
                )
                rename_btn = gr.Button("💾", size="sm", scale=1)
            
            def rename_session_cb(new_title, selected_label):
                sid = label_to_id(selected_label)
                if sid and new_title.strip():
                    conn = db()
                    cur = conn.cursor()
                    cur.execute("UPDATE sessions SET title=? WHERE id=?", (new_title.strip(), sid))
                    conn.commit()
                    conn.close()
                labels, _ = list_sessions()
                new_selected = next((lbl for lbl in labels if lbl.startswith(f"{sid} ")), None)
                return gr.update(choices=labels, value=new_selected)
            
            rename_btn.click(rename_session_cb, inputs=[edit_title_box, session_list], outputs=session_list)

            # Model and Dataset in a compact row
            with gr.Row(elem_classes=["compact-row"]):
                with gr.Column(scale=1):
                    gr.Markdown("**🤖 Model**")
                    model_choice = gr.Dropdown(
                        choices=list(MODELS.keys()),
                        value=list(MODELS.keys())[0],
                        label="",
                        interactive=True,
                        container=False
                    )
                with gr.Column(scale=1):
                    gr.Markdown("**📚 Dataset**")
                    dataset_choice = gr.Dropdown(
                        choices=DATASETS,
                        value=DATASETS[0],
                        label="",
                        interactive=True,
                        container=False
                    )

            # Generation Settings - more compact
            gr.Markdown("### ⚙️ Settings", elem_classes=["section-header"])
            
            # Compact system message
            system_box = gr.Textbox(
                value=SYSTEM_DEFAULT,
                label="System prompt",
                lines=2,
                max_lines=2,
                elem_classes=["small-textbox"]
            )
            
            # Sliders in compact layout
            max_tokens = gr.Slider(
                256, 2048, value=1200, step=16, 
                label="Max tokens",
                elem_classes=["compact-slider"]
            )
            temperature = gr.Slider(
                0.0, 2.0, value=0.25, step=0.05,
                label="Temperature", 
                elem_classes=["compact-slider"]
            )
            top_p = gr.Slider(
                0.1, 1.0, value=0.9, step=0.05,
                label="Top-p",
                elem_classes=["compact-slider"]
            )

        # Main chat area
        with gr.Column(scale=3):
            chatbot = gr.Chatbot(
                label="Assistant", 
                height=520, 
                type="messages",
                show_copy_button=True
            )
            with gr.Row():
                user_box = gr.Textbox(
                    placeholder="Describe your CRUD/backend task…", 
                    lines=2,
                    max_lines=3,
                    scale=5,
                    container=False
                )
            with gr.Row():
                send_btn = gr.Button("Send ▶️", variant="primary", size="sm")
                regen_btn = gr.Button("Regenerate 🔁", variant="secondary", size="sm")

    # Event handlers (unchanged)
    refresh_btn.click(refresh_sessions_cb, outputs=session_list)
    new_btn.click(new_chat_cb, outputs=[session_list, chatbot, user_box])
    del_btn.click(delete_chat_cb, inputs=session_list, outputs=[session_list, chatbot])
    session_list.change(load_session_cb, inputs=session_list, outputs=chatbot)

    send_btn.click(
        send_cb,
        inputs=[user_box, session_list, chatbot, system_box, max_tokens, temperature, top_p, model_choice, dataset_choice],
        outputs=[chatbot, user_box, session_list]
    )

    user_box.submit(
        send_cb,
        inputs=[user_box, session_list, chatbot, system_box, max_tokens, temperature, top_p, model_choice, dataset_choice],
        outputs=[chatbot, user_box, session_list]
    )

    regen_btn.click(
        regenerate_cb,
        inputs=[session_list, system_box, max_tokens, temperature, top_p, model_choice, dataset_choice],
        outputs=chatbot
    )

if __name__ == "__main__":
    demo.launch()