#!/usr/bin/env python3
"""
Chat Local - Asistente de IA offline
------------------------------------
Usa Ollama para ejecutar un modelo de lenguaje en tu propio ordenador.
No necesita internet, no consume APIs, todo corre localmente.

Requisitos:
1. Instalar Ollama desde https://ollama.com
2. Ejecutar: ollama pull llama3.2:3b  (o el modelo que prefieras)
3. Ejecutar este script: python3 chat_local.py

Más modelos disponibles: ollama run llama3.2:1b (más rápido, menos preciso)
                          ollama run llama3.2:3b (recomendado)
                          ollama run mistral:7b (más potente, necesita más RAM)
"""

import subprocess
import sys
import os
import json
from datetime import datetime

HISTORIAL_FILE = "chat_historial.json"
MODELO = "llama3.2:3b"

def limpiar_pantalla():
    os.system('cls' if os.name == 'nt' else 'clear')

def mostrar_banner():
    print("""
╔══════════════════════════════════════╗
║       🤖 CHAT LOCAL - AI Distillery ║
║     Tu asistente IA sin internet     ║
║   Powered by Ollama + llama3.2:3b    ║
╚══════════════════════════════════════╝
    """)
    print(f"Modelo: {MODELO}")
    print(f"Escribe 'salir' para terminar, 'modelo' para cambiar modelo")
    print("-" * 50)

def verificar_ollama():
    try:
        result = subprocess.run(["ollama", "--version"], capture_output=True, text=True, timeout=5)
        return True
    except (FileNotFoundError, subprocess.TimeoutExpired):
        return False

def verificar_modelo():
    result = subprocess.run(["ollama", "list"], capture_output=True, text=True, timeout=10)
    return MODELO in result.stdout

def consultar_ollama(prompt, contexto=None):
    messages = []
    if contexto:
        for msg in contexto[-10:]:  # Últimos 10 mensajes de contexto
            messages.append(msg)
    messages.append({"role": "user", "content": prompt})

    payload = {
        "model": MODELO,
        "messages": messages,
        "stream": True
    }

    cmd = ["ollama", "run", MODELO, prompt]
    
    try:
        proceso = subprocess.Popen(
            cmd,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
            bufsize=1
        )
        
        respuesta_completa = ""
        for linea in proceso.stdout:
            sys.stdout.write(linea)
            sys.stdout.flush()
            respuesta_completa += linea
        
        proceso.wait()
        return respuesta_completa.strip()
    except Exception as e:
        return f"Error: {e}"

def guardar_historial(mensajes):
    with open(HISTORIAL_FILE, "w", encoding="utf-8") as f:
        json.dump(mensajes, f, ensure_ascii=False, indent=2)

def cargar_historial():
    if os.path.exists(HISTORIAL_FILE):
        with open(HISTORIAL_FILE, "r", encoding="utf-8") as f:
            return json.load(f)
    return []

def main():
    limpiar_pantalla()
    mostrar_banner()

    if not verificar_ollama():
        print("❌ Ollama no está instalado.")
        print("Descárgalo desde: https://ollama.com")
        print("Después de instalarlo, ejecuta: ollama pull llama3.2:3b")
        input("\nPresiona Enter para salir...")
        return

    if not verificar_modelo():
        print(f"⚠️  El modelo {MODELO} no está descargado.")
        print(f"Ejecuta: ollama pull {MODELO}")
        descargar = input("¿Quieres descargarlo ahora? (s/n): ")
        if descargar.lower() == 's':
            print(f"Descargando {MODELO}...")
            subprocess.run(["ollama", "pull", MODELO])
            print("✅ Descarga completada")
        else:
            print("Puedes cambiar el modelo con el comando 'modelo'")
    
    historial = cargar_historial()
    if historial:
        print(f"📂 Historial cargado ({len(historial)} mensajes)")
    
    print("\n💬 Puedes empezar a hablar con la IA. Escribe 'ayuda' para ver comandos.\n")

    while True:
        try:
            usuario = input("\n🧑 Tú: ").strip()
            
            if usuario.lower() == "salir":
                guardar_historial(historial)
                print("¡Hasta luego! 🚀")
                break
            elif usuario.lower() == "limpiar":
                limpiar_pantalla()
                mostrar_banner()
                continue
            elif usuario.lower() == "modelo":
                print("Modelos disponibles:")
                subprocess.run(["ollama", "list"])
                nuevo = input("Nuevo modelo (ej: llama3.2:1b): ").strip()
                if nuevo:
                    global MODELO
                    MODELO = nuevo
                    print(f"✅ Modelo cambiado a {MODELO}")
                continue
            elif usuario.lower() == "ayuda":
                print("\nComandos disponibles:")
                print("  salir   - Terminar la conversación")
                print("  limpiar - Limpiar la pantalla")
                print("  modelo  - Cambiar de modelo")
                print("  ayuda   - Mostrar esta ayuda")
                print("  nuevo   - Empezar una conversación nueva")
                continue
            elif usuario.lower() == "nuevo":
                historial = []
                print("🗑️  Conversación nueva. Historial borrado.")
                continue
            elif not usuario:
                continue

            historial.append({"role": "user", "content": usuario})
            
            print("\n🤖 IA: ", end="", flush=True)
            respuesta = consultar_ollama(usuario, historial)
            
            historial.append({"role": "assistant", "content": respuesta})
            
            # Guardar cada pocos mensajes
            if len(historial) % 5 == 0:
                guardar_historial(historial)
                
        except KeyboardInterrupt:
            print("\n\n👋 Interrupción detectada. Guardando historial...")
            guardar_historial(historial)
            break
        except Exception as e:
            print(f"\n❌ Error inesperado: {e}")

if __name__ == "__main__":
    main()