import os
import json
import joblib
import pandas as pd
from fastapi import FastAPI, Form, HTTPException
import httpx # A modern library for making API calls

# --- 1. Basic Setup & Configuration ---
app = FastAPI(title="Alysium Corporation Studios's Hybrid Auto-Training AI")

# The permanent ID for your AI, as you requested.
MASTER_AI_ID = "neurones_self"
USER_MODELS_DIR = "user_models_data"
os.makedirs(USER_MODELS_DIR, exist_ok=True)

# --- 2. Helper Functions ---
def get_ai_paths(ai_id: str = MASTER_AI_ID):
    """Gets the file paths for your master AI."""
    ai_dir = os.path.join(USER_MODELS_DIR, ai_id)
    os.makedirs(ai_dir, exist_ok=True)
    return {
        "model_path": os.path.join(ai_dir, "matcher_model.joblib"),
        "data_path": os.path.join(ai_dir, "training_pairs.csv"),
        "responses_path": os.path.join(ai_dir, "responses.json")
    }

async def train_local_ai(prompt: str, reply: str):
    """This function contains the logic to train your personal AI."""
    paths = get_ai_paths()
    
    # Manage the list of unique replies
    if os.path.exists(paths["responses_path"]):
        with open(paths["responses_path"], 'r') as f:
            responses = json.load(f)
    else:
        responses = []

    if reply not in responses:
        responses.append(reply)
        with open(paths["responses_path"], 'w') as f:
            json.dump(responses, f)
    
    reply_index = responses.index(reply)

    # Save the new training pair
    new_data = pd.DataFrame([{"prompt": prompt, "label": reply_index}])
    if os.path.exists(paths["data_path"]):
        new_data.to_csv(paths["data_path"], mode='a', header=False, index=False)
    else:
        new_data.to_csv(paths["data_path"], mode='w', header=True, index=False)
            
    # Retrain the AI model
    df = pd.read_csv(paths["data_path"])
    
    # The model needs at least two different examples to learn anything.
    if len(df['label'].unique()) < 2:
        return # Exit if we don't have enough data to train

    X = df['prompt']
    y = df['label']
    
    model_pipeline = Pipeline([
        ('tfidf', TfidfVectorizer()),
        ('clf', SGDClassifier(loss='hinge', penalty='l2', alpha=1e-3, random_state=42, max_iter=100, tol=None)),
    ])
    
    model_pipeline.fit(X, y)
    joblib.dump(model_pipeline, paths["model_path"])

async def get_generative_reply(prompt: str):
    """Gets a reply from the powerful external Generative AI."""
    system_prompt = "You are a helpful AI assistant. Be friendly, creative, and concise."
    final_prompt = f"{system_prompt}\n\nUser message: \"{prompt}\""
    api_url = "https://main-gemini-2-0-flash-large-language-model-j7a2x36pcq-uc.a.run.app"

    try:
        async with httpx.AsyncClient(timeout=60.0) as client:
            response = await client.post(api_url, json={"contents": [{"parts": [{"text": final_prompt}]}]})
            response.raise_for_status()
            result = response.json()
        
        if result.get("candidates"):
            return result["candidates"][0]["content"]["parts"][0]["text"].strip()
        else:
            return None
    except Exception:
        return None


# --- 3. API Endpoints ---

@app.get("/")
def read_root():
    return {"message": "Welcome! This is NeuraPrompt AI. It learns from every conversation."}

@app.post("/chat/")
async def chat(text: str = Form(...)):
    """The main chat endpoint with the hybrid auto-training logic."""
    paths = get_ai_paths()
    
    # --- Step 1: Check if YOUR AI already knows a confident answer ---
    if os.path.exists(paths["model_path"]):
        model_pipeline = joblib.load(paths["model_path"])
        with open(paths["responses_path"], 'r') as f:
            responses = json.load(f)
            
        probabilities = model_pipeline.predict_proba([text])[0]
        max_confidence = max(probabilities)
        
        # If confidence is very high, use the learned reply.
        if max_confidence > 0.95:
            predicted_index = probabilities.argmax()
            return {"reply": responses[predicted_index], "source": "neurones_self"}

    # --- Step 2: If not, get a new reply from the powerful Generative AI ---
    generative_reply = await get_generative_reply(text)
    
    if generative_reply:
        # --- Step 3: THE MAGIC - Automatically train your AI with the new knowledge ---
        await train_local_ai(prompt=text, reply=generative_reply)
        return {"reply": generative_reply, "source": "generative_ai"}
    else:
        raise HTTPException(status_code=503, detail="The generative AI service is currently unavailable.")

@app.post("/manual_train/")
async def manual_train(prompt: str = Form(...), reply: str = Form(...)):
    """A separate endpoint to manually teach your AI specific replies."""
    await train_local_ai(prompt=prompt, reply=reply)
    return {"message": "Manual training successful. neurones_self has learned a new reply."}