modules/text_ui.py

"""
TextAI UI Module
Clean ChatGPT/Grok style interface
"""
import json
from typing import List, Tuple, Optional
from datetime import datetime

from .text_ai import (
    model_manager, session_manager,
    DEFAULT_CHAT_PROMPT, DEFAULT_ROLEPLAY_PROMPT, ROLEPLAY_PRESETS
)
from .config import MODELS_DIR

# Model configurations storage (per-model system prompts)
MODEL_CONFIGS_FILE = MODELS_DIR / "model_configs.json"


def _load_model_configs() -> dict:
    """Load per-model configurations"""
    if MODEL_CONFIGS_FILE.exists():
        try:
            return json.loads(MODEL_CONFIGS_FILE.read_text())
        except:
            pass
    return {}


def _save_model_configs(configs: dict):
    """Save per-model configurations"""
    MODEL_CONFIGS_FILE.parent.mkdir(parents=True, exist_ok=True)
    MODEL_CONFIGS_FILE.write_text(json.dumps(configs, indent=2))


# ══════════════════════════════════════════════════════════════════════════════
# CHAT SESSION HELPERS
# ══════════════════════════════════════════════════════════════════════════════

def get_chat_list() -> List[List[str]]:
    """Get all chat sessions sorted by recent - clean display"""
    sessions = session_manager.list_sessions()  # All types
    rows = []
    for s in sessions[:30]:
        # Show only title (auto-generated or renamed)
        title = s["title"][:35] + "..." if len(s["title"]) > 35 else s["title"]
        session_type = "🎭" if s.get("session_type") == "roleplay" else "💬"
        rows.append([s["session_id"], f"{session_type} {title}"])
    return rows


def get_model_choices() -> List[str]:
    """Get available models for dropdown"""
    models = model_manager.get_available_models()
    choices = []
    for m in models:
        status = "✓ " if m["loaded"] else ""
        choices.append(f"{status}{m['name']} ({m['type']})")
    return choices if choices else ["No models available"]


def get_current_model_display() -> str:
    """Get current model for display"""
    status = model_manager.get_status()
    if status["loaded"]:
        return f"{status['model_id']}"
    return "No model loaded"


def format_chat_history(session_id: str) -> List[dict]:
    """Format session messages for Gradio Chatbot (new format)"""
    session = session_manager.load_session(session_id)
    if not session:
        return []

    history = []
    for msg in session.messages:
        if msg["role"] == "user":
            history.append({"role": "user", "content": msg["content"]})
        elif msg["role"] == "assistant":
            history.append({"role": "assistant", "content": msg["content"]})

    return history


def format_chat_history_tuples(session_id: str) -> List[Tuple[str, str]]:
    """Format session messages for Gradio Chatbot (tuple format for compatibility)"""
    session = session_manager.load_session(session_id)
    if not session:
        return []

    history = []
    user_msg = None

    for msg in session.messages:
        if msg["role"] == "user":
            user_msg = msg["content"]
        elif msg["role"] == "assistant" and user_msg:
            history.append((user_msg, msg["content"]))
            user_msg = None

    if user_msg:
        history.append((user_msg, None))

    return history


# ══════════════════════════════════════════════════════════════════════════════
# MAIN CHAT FUNCTIONS
# ══════════════════════════════════════════════════════════════════════════════

def ui_new_chat(mode: str = "chat"):
    """Create new chat session"""
    # Get system prompt based on mode and current model
    if mode == "roleplay":
        system_prompt = DEFAULT_ROLEPLAY_PROMPT
        session_type = "roleplay"
    else:
        # Use model-specific prompt if available
        model_id = model_manager.current_model_id
        configs = _load_model_configs()
        system_prompt = configs.get(model_id, {}).get("system_prompt", DEFAULT_CHAT_PROMPT)
        session_type = "chat"

    session = session_manager.create_session("", session_type, system_prompt)
    return (
        session.session_id,
        [],
        get_chat_list(),
        session.title
    )


def ui_load_session(evt, sessions_data):
    """Load session from sidebar click"""
    try:
        if hasattr(evt, 'index') and evt.index[0] < len(sessions_data):
            session_id = sessions_data[evt.index[0]][0]
            session = session_manager.load_session(session_id)
            if session:
                history = format_chat_history_tuples(session_id)
                return session_id, history, session.title, session.session_type == "roleplay"
    except:
        pass
    return "", [], "", False


def ui_send_message(
    session_id: str,
    message: str,
    history: List,
    max_tokens: int,
    temperature: float,
    is_roleplay: bool = False
):
    """Send message and stream response"""
    if not session_id:
        # Auto-create session
        mode = "roleplay" if is_roleplay else "chat"
        session_id, _, _, _ = ui_new_chat(mode)

    if not message.strip():
        yield history, "", session_id, get_chat_list()
        return

    if model_manager.current_model is None:
        history = history + [(message, "Please load a model first from the menu.")]
        yield history, "", session_id, get_chat_list()
        return

    # Add user message
    history = history + [(message, None)]
    yield history, "", session_id, get_chat_list()

    # Get response
    from .text_ai import api_chat
    result = json.loads(api_chat(session_id, message, max_tokens, temperature))

    if result["success"]:
        history[-1] = (message, result["response"])
        # Update title if first message
        session = session_manager.load_session(session_id)
        title = session.title if session else ""
    else:
        history[-1] = (message, f"Error: {result.get('error', 'Unknown error')}")
        title = ""

    yield history, "", session_id, get_chat_list()


def ui_rename_session(session_id: str, new_title: str):
    """Rename current session"""
    if session_id and new_title.strip():
        session_manager.rename_session(session_id, new_title.strip())
        return get_chat_list(), new_title.strip()
    return get_chat_list(), ""


def ui_delete_session(session_id: str):
    """Delete current session"""
    if session_id:
        session_manager.delete_session(session_id)
    return "", [], get_chat_list(), ""


def ui_clear_chat(session_id: str):
    """Clear current chat messages"""
    if session_id:
        session_manager.clear_session(session_id)
    return []


# ══════════════════════════════════════════════════════════════════════════════
# MODEL MANAGEMENT (for Tools tab)
# ══════════════════════════════════════════════════════════════════════════════

def get_models_table() -> List[List[str]]:
    """Get models for table display"""
    models = model_manager.get_available_models()
    rows = []
    for m in models:
        configs = _load_model_configs()
        has_prompt = "✓" if m["id"] in configs else ""
        rows.append([
            "●" if m["loaded"] else "",
            m["name"],
            m["type"],
            m["size"],
            has_prompt
        ])
    return rows


def ui_load_model_by_index(evt, models_data):
    """Load model by clicking on table row"""
    try:
        if hasattr(evt, 'index') and evt.index[0] < len(models_data):
            models = model_manager.get_available_models()
            if evt.index[0] < len(models):
                model_id = models[evt.index[0]]["id"]
                result = model_manager.load_model(model_id)
                status = f"Loaded: {model_id}" if result.get("success") else f"Error: {result.get('error')}"
                return get_models_table(), get_current_model_display(), status
    except Exception as e:
        return get_models_table(), get_current_model_display(), f"Error: {str(e)}"
    return get_models_table(), get_current_model_display(), ""


def ui_unload_model():
    """Unload current model"""
    model_manager.unload_model()
    return get_models_table(), get_current_model_display(), "Model unloaded"


def ui_save_model_prompt(model_name: str, system_prompt: str):
    """Save system prompt for a model"""
    if not model_name:
        return "Select a model first"

    # Find model ID from name
    models = model_manager.get_available_models()
    model_id = None
    for m in models:
        if m["name"] == model_name or m["id"] == model_name:
            model_id = m["id"]
            break

    if not model_id:
        return "Model not found"

    configs = _load_model_configs()
    if model_id not in configs:
        configs[model_id] = {}
    configs[model_id]["system_prompt"] = system_prompt
    _save_model_configs(configs)

    return f"Saved prompt for {model_name}"


def ui_get_model_prompt(model_name: str) -> str:
    """Get system prompt for a model"""
    models = model_manager.get_available_models()
    model_id = None
    for m in models:
        if m["name"] == model_name or m["id"] == model_name:
            model_id = m["id"]
            break

    if model_id:
        configs = _load_model_configs()
        return configs.get(model_id, {}).get("system_prompt", DEFAULT_CHAT_PROMPT)
    return DEFAULT_CHAT_PROMPT


def ui_delete_model_config(model_name: str):
    """Delete model from config (not the file)"""
    models = model_manager.get_available_models()
    model_id = None
    for m in models:
        if m["name"] == model_name:
            model_id = m["id"]
            break

    if model_id:
        configs = _load_model_configs()
        if model_id in configs:
            del configs[model_id]
            _save_model_configs(configs)
            return get_models_table(), "Config removed"
    return get_models_table(), "Model not found"


# ══════════════════════════════════════════════════════════════════════════════
# PERSONA MANAGEMENT
# ══════════════════════════════════════════════════════════════════════════════

PERSONAS_FILE = MODELS_DIR.parent / "storage" / "personas.json"


def _load_personas() -> dict:
    """Load saved personas"""
    if PERSONAS_FILE.exists():
        try:
            return json.loads(PERSONAS_FILE.read_text())
        except:
            pass
    return {"default": {"name": "Default", "prompt": DEFAULT_CHAT_PROMPT}}


def _save_personas(personas: dict):
    """Save personas"""
    PERSONAS_FILE.parent.mkdir(parents=True, exist_ok=True)
    PERSONAS_FILE.write_text(json.dumps(personas, indent=2))


def get_persona_choices() -> List[str]:
    """Get persona names for dropdown"""
    personas = _load_personas()
    return list(personas.keys())


def ui_save_persona(name: str, prompt: str):
    """Save a persona"""
    if not name.strip():
        return get_persona_choices(), "Enter persona name"

    personas = _load_personas()
    personas[name.strip()] = {"name": name.strip(), "prompt": prompt}
    _save_personas(personas)
    return get_persona_choices(), f"Saved: {name}"


def ui_load_persona(name: str) -> str:
    """Load persona prompt"""
    personas = _load_personas()
    return personas.get(name, {}).get("prompt", DEFAULT_CHAT_PROMPT)


def ui_delete_persona(name: str):
    """Delete a persona"""
    if name == "default":
        return get_persona_choices(), "Cannot delete default persona"

    personas = _load_personas()
    if name in personas:
        del personas[name]
        _save_personas(personas)
        return get_persona_choices(), f"Deleted: {name}"
    return get_persona_choices(), "Persona not found"


# ══════════════════════════════════════════════════════════════════════════════
# SUGGESTED MODELS
# ══════════════════════════════════════════════════════════════════════════════

# Default model for quick testing (small, fast, uncensored)
DEFAULT_MODEL = {
    "id": "TheBloke/TinyLlama-1.1B-Chat-v1.0-GGUF",
    "file": "tinyllama-1.1b-chat-v1.0.Q4_K_M.gguf",
    "name": "TinyLlama 1.1B",
    "size": "0.7GB"
}

SUGGESTED_MODELS = [
    {"id": "TheBloke/TinyLlama-1.1B-Chat-v1.0-GGUF", "file": "tinyllama-1.1b-chat-v1.0.Q4_K_M.gguf", "name": "TinyLlama 1.1B (Fast)", "size": "0.7GB", "recommended": True},
    {"id": "TheBloke/phi-2-GGUF", "file": "phi-2.Q4_K_M.gguf", "name": "Phi-2 (Small & Fast)", "size": "1.6GB"},
    {"id": "Qwen/Qwen2-0.5B-Instruct-GGUF", "file": "qwen2-0_5b-instruct-q4_k_m.gguf", "name": "Qwen2 0.5B (Tiny)", "size": "0.4GB"},
    {"id": "TheBloke/Mistral-7B-Instruct-v0.2-GGUF", "file": "mistral-7b-instruct-v0.2.Q4_K_M.gguf", "name": "Mistral 7B Instruct", "size": "4.4GB"},
    {"id": "TheBloke/Llama-2-7B-Chat-GGUF", "file": "llama-2-7b-chat.Q4_K_M.gguf", "name": "Llama 2 7B Chat", "size": "4.1GB"},
    {"id": "TheBloke/OpenHermes-2.5-Mistral-7B-GGUF", "file": "openhermes-2.5-mistral-7b.Q4_K_M.gguf", "name": "OpenHermes 2.5", "size": "4.4GB"},
]


def download_default_model() -> str:
    """Download the default model if not present"""
    from .hf_hub import download_model_file
    from .config import MODELS_DIR

    txt_dir = MODELS_DIR / "txt"
    txt_dir.mkdir(parents=True, exist_ok=True)

    # Check if any model exists
    existing = list(txt_dir.rglob("*.gguf"))
    if existing:
        return f"Model already exists: {existing[0].name}"

    # Download default model
    result = download_model_file(DEFAULT_MODEL["id"], DEFAULT_MODEL["file"])
    return result


def ensure_model_available() -> bool:
    """Ensure at least one model is available, download if needed"""
    models = model_manager.get_available_models()
    return len(models) > 0


def get_suggested_models_table() -> List[List[str]]:
    """Get suggested models for display"""
    return [[m["name"], m["size"], m["id"]] for m in SUGGESTED_MODELS]


# ══════════════════════════════════════════════════════════════════════════════
# COMPATIBILITY EXPORTS
# ══════════════════════════════════════════════════════════════════════════════

def get_sessions_list(session_type: str = None) -> List[List[str]]:
    """Legacy: Get sessions list"""
    return get_chat_list()


def get_model_status_display() -> str:
    """Legacy: Get model status"""
    status = model_manager.get_status()
    if status["loaded"]:
        return f"Model: {status['model_id']}"
    return "No model loaded"


def ui_refresh_models():
    """Legacy: Refresh models"""
    return get_models_table(), get_current_model_display()


def ui_load_model(evt, models_data):
    """Legacy: Load model"""
    return ui_load_model_by_index(evt, models_data)