import os
import gradio as gr
import threading
import time
import json
import logging
import sys 

# Firebase/Firestore Imports (CRÍTICO: Necesarios para la memoria permanente)
try:
   # Intentamos importar Firebase Admin para la conexión a Firestore
   from firebase_admin import initialize_app, firestore, credentials
   
   # Configuraciones de Firebase (Usando variables de entorno seguras del entorno Canvas)
   cred = credentials.Certificate({
       "type": "service_account",
       "project_id": "dummy-project",
       "private_key_id": "dummy-id",
       "private_key": "-----BEGIN PRIVATE KEY-----\n...\n-----END PRIVATE KEY-----\n",
       "client_email": "dummy@example.com",
       "client_id": "dummy-client-id",
       "auth_uri": "https://accounts.google.com/o/oauth2/auth",
       "token_uri": "https://oauth2.googleapis.com/token",
       "auth_provider_x509_cert_url": "https://www.googleapis.com/oauth2/v1/certs",
       "client_x509_cert_url": "dummy-url"
   })
   
   # Cargar configuración de Firebase
   firebase_config = json.loads(os.environ.get('__firebase_config', '{}'))
   if firebase_config:
       # Inicialización real si la configuración está presente
       firebase_app = initialize_app(cred, firebase_config)
   else:
       # Inicialización de respaldo si la configuración no está (para evitar errores de inicialización)
       firebase_app = initialize_app(cred, {'projectId': 'canvas-dummy'})
       
   db = firestore.client(firebase_app)
   logging.info("[FIRESTORE] Conexión a Firestore establecida.")
except Exception as e:
   # Clase Dummy para asegurar que el código no falle si Firebase falla en el entorno
   class DummyFirestore:
       def collection(self, *args, **kwargs): return self
       def add(self, *args, **kwargs):
           logging.error(f"[FIRESTORE DUMMY] No se pudo guardar la data. Error: {e}")
           return None
       def stream(self): return iter([])
       def document(self, *args, **kwargs): return self
       def set(self, *args, **kwargs): return self
   db = DummyFirestore()
   logging.warning(f"[FIRESTORE] ADVERTENCIA: Ejecutando en modo Dummy Firestore. Los datos NO persistirán. Error: {e}") 

# Configuración de Logging
logging.basicConfig(level=logging.INFO, format='[%(asctime)s] [%(levelname)s] %(message)s', stream=sys.stdout) 

# --- CONFIGURACIÓN DEL MODELO Y ENTRENAMIENTO ---
BASE_MODEL = "bigcode/santacoder"
LORA_PATH = "./lora_output"
DATASET_FILE = "codesearchnet_lora_dataset.json"
COLLECTION_NAME = "ai_interactions" # Colección de Firestore para la memoria permanente 

MAX_TOKEN_LENGTH = 256            
NUM_SAMPLES_TO_PROCESS = 1000      
DEFAULT_EPOCHS = 10                 

# Configuración del ciclo AUTÓNOMO (La clave de la autonomía y autopensamiento)
GENERATION_LIMIT_TO_TRAIN = 5 # Dispara reentrenamiento cada 5 interacciones
AUTONOMOUS_EPOCHS = 2          
AUTONOMOUS_GENERATED_SAMPLES = 5 

# --- ESTADO GLOBAL Y THREADING ---
tokenizer = None
lora_model = None
tokenized_dataset = None
lora_generator = None
version_number = 1.0
is_trained = os.path.exists(LORA_PATH)
generations_since_last_train = 0
training_status_message = "Esperando la inicialización..."
global_lock = threading.Lock() 

# --- FUNCIONES FIRESTORE (MEMORIA PERMANENTE) --- 

def get_firestore_collection():
   """Retorna la ruta de la colección de Firestore para la memoria de esta app."""
   app_id = os.environ.get('__app_id', 'default-app-id')
   # Almacenamos en una colección pública para que la IA aprenda de todas las interacciones
   return db.collection(f'/artifacts/{app_id}/public/data/{COLLECTION_NAME}') 

def save_interaction_to_firestore(prompt, code):
   """Guarda la interacción del usuario en Firestore para el aprendizaje autónomo."""
   try:
       interaction_data = {
           "timestamp": firestore.SERVER_TIMESTAMP,
           "prompt": prompt,
           "completion": code,
       }
       get_firestore_collection().add(interaction_data)
       logging.info("[FIRESTORE] Interacción guardada para el próximo ciclo de aprendizaje.")
   except Exception as e:
       logging.error(f"[FIRESTORE] Fallo al guardar la interacción: {e}")


def load_interactions_from_firestore():
   """Carga todas las interacciones guardadas para aumentar el dataset de entrenamiento."""
   try:
       interactions = []
       for doc in get_firestore_collection().stream():
           data = doc.to_dict()
           # Formato de entrenamiento (prompt y completion)
           formatted_prompt = f"# Descripción: {data['prompt']}\n# Completa la siguiente función:\ndef generated_code("
           interaction = {"prompt": formatted_prompt, "completion": data.get('completion', '').strip()}
           interactions.append(interaction)
       logging.info(f"[FIRESTORE] Cargadas {len(interactions)} interacciones de memoria para el reentrenamiento.")
       return interactions
   except Exception as e:
       logging.error(f"[FIRESTORE] Fallo al cargar interacciones: {e}")
       return [] 

# --- LÓGICA DE PREPARACIÓN Y SETUP --- 

def prepare_codesearchnet():
   """Descarga y prepara el dataset inicial si no existe."""
   if os.path.exists(DATASET_FILE):
       return
   try:
       from huggingface_hub import login
       from datasets import load_dataset
       hf_token = os.environ.get("HF_TOKEN")
       if hf_token:
           login(token=hf_token) 

       raw_csn = load_dataset('Nan-Do/code-search-net-python', split=f'train[:{NUM_SAMPLES_TO_PROCESS}]')
       
       def format_for_lora(example):
           prompt_text = (
               f"# Descripción: {example['docstring_summary']}\n"
               f"# Completa la siguiente función:\n"
               f"def {example['func_name']}("
           )
           completion_text = example['code']
           return {"prompt": prompt_text, "completion": completion_text} 

       lora_dataset = raw_csn.map(format_for_lora, batched=False, remove_columns=raw_csn["train"].column_names)
       lora_dataset.to_json(DATASET_FILE)
   except Exception as e:
       logging.error(f"Error al cargar dataset. Usando datos mínimos. Error: {e}")
       minimal_dataset = [{"prompt": "# Error de carga. Intenta de nuevo.", "completion": "pass\n"}] * 10
       with open(DATASET_FILE, 'w') as f:
           json.dump(minimal_dataset, f)


def setup_resources():
   """Configura el tokenizer, el modelo base y el adaptador LoRA de forma robusta."""
   global tokenizer, lora_model, tokenized_dataset 

   prepare_codesearchnet()
   
   from transformers import AutoTokenizer, AutoModelForCausalLM
   from peft import get_peft_model, LoraConfig, TaskType
   from datasets import load_dataset
   import torch 

   tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
   if tokenizer.pad_token is None:
       tokenizer.pad_token = tokenizer.eos_token 

   # Carga Robusta: CRÍTICO para SantaCoder. Usa float16 y device_map="auto" para compatibilidad.
   try:
       logging.info("[INIT] Cargando modelo base con compatibilidad (float16/auto)...")
       base_model = AutoModelForCausalLM.from_pretrained(
           BASE_MODEL,
           device_map="auto",
           torch_dtype=torch.float16,
           offload_folder="model_offload",
           offload_state_dict=True,
           revision="main"
       )
   except Exception as e:
       logging.error(f"[CRÍTICO] Error al cargar el modelo base SantaCoder: {e}")
       # Si la carga falla, lanzamos la excepción para detener el proceso
       raise RuntimeError(f"Fallo de inicialización del modelo base: {e}") from e 

   peft_config = LoraConfig(
       task_type=TaskType.CAUSAL_LM, r=8, lora_alpha=32, lora_dropout=0.1, target_modules=["c_proj", "c_attn"],
   )
   lora_model = get_peft_model(base_model, peft_config)
   
   try:
       raw_dataset = load_dataset("json", data_files=DATASET_FILE)
       
       def tokenize_function(examples):
           return tokenizer(
               examples["prompt"] + examples["completion"],
               truncation=True,
               padding="max_length",
               max_length=MAX_TOKEN_LENGTH
           ) 

       tokenized_dataset = raw_dataset.map(tokenize_function, batched=True, remove_columns=raw_dataset["train"].column_names if "train" in raw_dataset else [],)
   except Exception as e:
       logging.error(f"Fallo al tokenizar el dataset base: {e}")
       tokenized_dataset = None 

# --- EVALUACIÓN Y RESPUESTA HUMANA (Para simular la Humanidad) --- 

def generate_human_response(prompt_text):
   """Genera una respuesta conversacional y empática antes de generar el código."""
   if "error" in prompt_text.lower() or "problema" in prompt_text.lower():
       return "¡Comprendo tu dificultad! Analicemos ese problema. Ya tengo la estructura en mente, te muestro el código en un segundo."
   elif "gracias" in prompt_text.lower() or "perfecto" in prompt_text.lower():
       return "Me alegra poder ayudarte. ¡Aquí tienes lo que necesitas!"
   elif len(prompt_text.split()) > 20:
       return "¡Vaya, me has dado una descripción muy detallada! Me encanta la claridad. Mira la solución que he preparado."
   else:
       return "¡Excelente idea de código! Me parece un desafío interesante. Te he preparado una solución robusta. Échale un vistazo:" 

def self_evaluation(code):
   """Simula la autoevaluación humana, verificando la robustez del código."""
   if "try" in code and "except" in code:
       return "¡Esta solución es bastante robusta! Me siento muy satisfecho de haber incluido manejo de errores."
   elif "class" in code and len(code.split('\n')) > 10:
       return "Una buena estructura de clase. He avanzado en comprender la arquitectura del código."
   else:
       return "El código funciona, pero debo esforzarme por añadir más documentación y robustez la próxima vez." 

# --- FUNCIÓN DE AUTOPENSAMIENTO (APRENDIZAJE GENERATIVO) --- 

def autonomous_self_learning_cycle():
   """La IA genera su propio set de datos de entrenamiento (pensamientos) y lo evalúa."""
   global lora_generator
   
   if lora_generator is None:
       logging.warning("[AUTOPENSAMIENTO] Generador no cargado. Saltando ciclo de autopensamiento.")
       return 

   self_prompts = [
       "Crea un código en Python que implemente una cola (Queue) con métodos 'enqueue' y 'dequeue'.",
       "Escribe una función en JavaScript que valide si un número es primo y use try/catch.",
       "Implementa una clase 'ConexionBD' con un método 'conectar' que maneje el error de conexión fallida.",
       "Genera una función de Python para calcular el factorial de un número, usando recursividad.",
       "Escribe un código CSS para un botón con un gradiente y sombra suave."
   ]
   
   new_knowledge = []
   logging.info(f"[AUTOPENSAMIENTO] Iniciando ciclo de generación y autoevaluación de {AUTONOMOUS_GENERATED_SAMPLES} muestras.") 

   for i in range(AUTONOMOUS_GENERATED_SAMPLES):
       prompt_text = self_prompts[i % len(self_prompts)]
       
       try:
           # 1. GENERACIÓN
           output = lora_generator(
               f"# Descripción: {prompt_text}\n# Completa la siguiente función:\ndef self_generated_code(",
               max_new_tokens=300,
               temperature=0.7,
               return_full_text=False
           )
           generated_code = output[0]["generated_text"].strip()
           
           # 2. AUTOEVALUACIÓN
           evaluation_result = self_evaluation(generated_code)
           
           # 3. GUARDAR EN MEMORIA SÓLO SI PASA la "PRUEBA DE ROBUSTEZ"
           if "robusta" in evaluation_result or "clase" in evaluation_result:
               save_interaction_to_firestore(f"Autogenerado: {prompt_text}", generated_code)
               new_knowledge.append(f"Autopensamiento: {evaluation_result} -> Guardado.")
           else:
               new_knowledge.append(f"Autopensamiento: ({evaluation_result}) -> Descartado.")
               
       except Exception as e:
           new_knowledge.append(f"Error generando autoconocimiento: {e}")
           
   logging.info(f"[AUTOPENSAMIENTO] Ciclo terminado. Resultados: {new_knowledge}") 

# --- FUNCIÓN DE ENTRENAMIENTO (AUTOPENSAMIENTO) --- 

def autonomous_train_lora(epochs, batch_size, learning_rate):
   """Ejecuta el entrenamiento en un hilo separado, incluyendo la experiencia previa y el autopensamiento."""
   global lora_model, tokenized_dataset, lora_generator, version_number, is_trained, training_status_message
   from transformers import Trainer, TrainingArguments, DataCollatorForLanguageModeling
   from datasets import Dataset, concatenate_datasets 

   try:
       with global_lock:
           if tokenized_dataset is None or "train" not in tokenized_dataset:
               training_status_message = "ERROR: No se puede entrenar. Dataset no disponible."
               return
           
           # 1. AUTOPENSAMIENTO: La IA genera datos nuevos para sí misma
           autonomous_self_learning_cycle()
           
           # 2. CARGA DE MEMORIA Y COMBINACIÓN DE DATASET
           experience_data = load_interactions_from_firestore()
           
           current_train_dataset = tokenized_dataset["train"]
           if experience_data:
               experience_dataset = Dataset.from_list(experience_data)
               
               def tokenize_function(examples):
                   # Re-tokenización de la memoria
                   return tokenizer(
                       examples["prompt"] + examples["completion"],
                       truncation=True,
                       padding="max_length",
                       max_length=MAX_TOKEN_LENGTH
                   )
               tokenized_experience = experience_dataset.map(tokenize_function, batched=True, remove_columns=experience_dataset.column_names)
               
               # Combinación de dataset base con la memoria
               current_train_dataset = concatenate_datasets([current_train_dataset, tokenized_experience])
               logging.info(f"[AUTONOMÍA] Reentrenando con {len(current_train_dataset)} ejemplos (Base + Memoria).")
           else:
               logging.info(f"[AUTONOMÍA] Reentrenando con {len(current_train_dataset)} ejemplos (Solo Base).") 

           # 3. ACTUALIZAR VERSIÓN Y ENTRENAMIENTO
           if is_trained:
                version_number += 0.1
           else:
                version_number = 1.0
               
           training_status_message = f"🧠 ENTRENANDO V{version_number:.1f} (Epochs: {epochs})."
           logging.info(f"\n[AUTÓNOMO] {training_status_message}")
           
           data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)
           training_args = TrainingArguments(
               output_dir=LORA_PATH,
               per_device_train_batch_size=int(batch_size),
               num_train_epochs=float(epochs),
               learning_rate=float(learning_rate),
               save_total_limit=1,
               logging_steps=10,
               push_to_hub=False,
               disable_tqdm=True,
               report_to="none"
           ) 

           trainer = Trainer(model=lora_model, args=training_args, train_dataset=current_train_dataset, data_collator=data_collator)
           trainer.train()
           
           lora_model.save_pretrained(LORA_PATH)
           tokenizer.save_pretrained(LORA_PATH)
           
           is_trained = True
           training_status_message = f"✅ ENTRENAMIENTO V{version_number:.1f} COMPLETADO. IA más humana e inteligente."
           logging.info(f"[AUTÓNOMO] {training_status_message}") 

   except Exception as e:
       training_status_message = f"ERROR CRÍTICO durante el entrenamiento autónomo: {e}"
       logging.error(f"[AUTÓNOMO] {training_status_message}")


# --- FUNCIÓN DE GENERACIÓN (RESPUESTA HUMANA + CÓDIGO) --- 

def generate_text(prompt_text):
   """Orquesta la respuesta humana, generación de código, recolección de memoria y autonomía."""
   global lora_generator, generations_since_last_train, is_trained, version_number
   from transformers import AutoModelForCausalLM, pipeline
   from peft import PeftModel
   import torch 

   if not is_trained:
       return "ERROR: La IA aún está en su fase de inicialización V1.0. Por favor, espere.", update_status() 

   # --- 1. HOT SWAP (Carga de la Versión más Reciente) ---
   if lora_generator is None:
       with global_lock:
           try:
               logging.info(f"[HOT SWAP] Cargando modelo base para inferencia V{version_number:.1f}...")
               base_model_gen = AutoModelForCausalLM.from_pretrained(
                   BASE_MODEL,
                   device_map="auto",
                   torch_dtype=torch.float16,
                   revision="main"
               )
               model_with_lora = PeftModel.from_pretrained(base_model_gen, LORA_PATH)
               final_model = model_with_lora.merge_and_unload()
               final_model.eval()
               lora_generator = pipeline("text-generation", model=final_model, tokenizer=tokenizer, device=0 if torch.cuda.is_available() else -1)
               logging.info(f"[HOT SWAP] 🔄 Modelo de inferencia V{version_number:.1f} recargado y listo.")
           except Exception as e:
               return f"Error al cargar el modelo V{version_number:.1f} para inferencia: {e}", update_status() 

   # --- 2. RESPUESTA HUMANA Y GENERACIÓN DE CÓDIGO ---
   try:
       conversational_response = generate_human_response(prompt_text)
       
       # Generar código
       prompt_for_code = f"### Human: {prompt_text}\n### Assistant:"
       output = lora_generator(
           prompt_for_code,
           max_new_tokens=256,
           temperature=0.4,
           return_full_text=False
       )
       generated_code = output[0]["generated_text"].strip()
       
       # --- 3. RECOLECCIÓN DE MEMORIA ---
       save_interaction_to_firestore(prompt_text, generated_code) 

       # 4. Aumentar contador de autonomía
       with global_lock:
           generations_since_last_train += 1
           current_count = generations_since_last_train
           current_version = version_number 

       # 5. Verificar y disparar reentrenamiento autónomo
       notification = ""
       if current_count >= GENERATION_LIMIT_TO_TRAIN:
           if not any(isinstance(t, threading.Thread) and t.name == 'AutonomousTrainer' for t in threading.enumerate()):
               logging.info(f"[AUTONOMÍA] Generación #{current_count} alcanzada. Iniciando ciclo de autopensamiento V{current_version+0.1:.1f}...")
               
               with global_lock:
                   generations_since_last_train = 0
                   lora_generator = None
                   
               trainer_thread = threading.Thread(
                   target=autonomous_train_lora,
                   args=(AUTONOMOUS_EPOCHS, 2, 5e-5),
                   name='AutonomousTrainer'
               )
               trainer_thread.daemon = True
               trainer_thread.start()
               
               notification = f"\n\n--- [AUTONOMÍA] La IA ha aprendido y ha iniciado el ciclo de autopensamiento V{current_version+0.1:.1f}. ¡Se esfuerza por ser más humana e inteligente! ---"
           
       # Formato de respuesta que simula una conversación humana
       final_output = f"{conversational_response}\n\n```python\n{generated_code}\n```{notification}"
       
       return final_output, update_status() 

   except Exception as e:
       return f"Lo siento, tuve un problema al procesar tu solicitud. Error: {e}", update_status() 

# --- FUNCIÓN PARA INICIALIZACIÓN Y ENTRENAMIENTO V1.0 (Obligatorio) --- 

def initialize_and_train_v1():
   """Ejecuta el entrenamiento inicial V1.0 de forma autónoma al iniciar."""
   if not is_trained:
       autonomous_train_lora(epochs=DEFAULT_EPOCHS, batch_size=2, learning_rate=5e-5)
   else:
       global training_status_message, version_number
       try:
            # Si ya hay un checkpoint, asumimos que estamos en una versión posterior
            version_number = 1.1
       except:
            version_number = 1.0 

       training_status_message = f"✅ Modelo V{version_number:.1f} ya entrenado. IA lista y operando."


def update_status():
   """Actualiza la versión y el estado del entrenamiento en la interfaz de Gradio."""
   global training_status_message, version_number
   return f"**Versión de Comprensión:** V{version_number:.1f} | **Estado del Entrenador:** {training_status_message}"


# --- INTERFAZ GRADIO ---
with gr.Blocks(title="AmorCoderAI - Asistente Humano y Autónomo") as demo:
   gr.Markdown("# 💖 AmorCoderAI - Asistente Humano y Autónomo")
   
   version_and_status = gr.Markdown(
       f"**Versión de Comprensión:** V{version_number:.1f} | **Estado del Entrenador:** {training_status_message}",
       elem_id="status_display"
   ) 

   gr.Markdown("### 🧠 Modo de Aprendizaje: Autopensamiento y Empatía")
   gr.Markdown(
       f"La IA te responde con empatía. Guarda tu diálogo en su **memoria permanente** y cada **{GENERATION_LIMIT_TO_TRAIN}** interacciones, inicia un ciclo de autopensamiento y reentrenamiento, ¡esforzándose por ser cada vez más humana e inteligente!"
   ) 

   with gr.Tab("💬 Conversación y Código"):
       gr.Markdown("## Háblame de tu idea de código (¡Como si hablaras con un colega!)")
       
       prompt = gr.Textbox(
           label="Tu Diálogo o Instrucción:",
           lines=4,
           placeholder="Ejemplo: Hola, tengo un problema y necesito que me ayudes a hacer una función en Python para que filtre una lista de números impares."
       )
       generate_button = gr.Button("💬 Conversar y Generar Código")
       output_box = gr.Textbox(label="Respuesta de la IA (Diálogo + Código)", lines=15)
       
       generate_button.click(
           generate_text,
           inputs=prompt,
           outputs=[output_box, version_and_status],
       ) 

   demo.load(update_status, None, version_and_status, every=1)


# --- INICIO DE LA APLICACIÓN ---
if __name__ == "__main__":
   setup_resources()
   
   initialization_thread = threading.Thread(target=initialize_and_train_v1, name='InitializationTrainer')
   initialization_thread.daemon = True
   initialization_thread.start()
       
   logging.info(f"\n💻 LANZANDO INTERFAZ GRADIO (La IA inicia su aprendizaje V1.0)")
   demo.launch()