import os
import json
import uuid
from dotenv import load_dotenv
from fasthtml.common import *
from openai import OpenAI
from groq import Groq
from google import genai 

# ==========================================
# SECTION 1. INITIALISATION, SÉCURITÉ & HISTORIQUE
# ==========================================
load_dotenv()

HISTORY_FILE = "aethas_history.json"
def load_history():
    if not os.path.exists(HISTORY_FILE):
        with open(HISTORY_FILE, 'w', encoding='utf-8') as f: json.dump({}, f)
    with open(HISTORY_FILE, 'r', encoding='utf-8') as f: return json.load(f)

def save_history(history):
    with open(HISTORY_FILE, 'w', encoding='utf-8') as f: json.dump(history, f, indent=4)

theme_hdrs = (
    Script(src="https://cdn.jsdelivr.net/npm/marked/marked.min.js"),
    Script(src="https://cdnjs.cloudflare.com/ajax/libs/highlight.js/11.9.0/highlight.min.js"),
    Script(src="https://cdnjs.cloudflare.com/ajax/libs/jszip/3.10.1/jszip.min.js"),
    Link(rel="stylesheet", href="https://cdnjs.cloudflare.com/ajax/libs/highlight.js/11.9.0/styles/github-dark.min.css"),
    
    Style("""
        :root { --bg-dark: #121212; --bg-panel: #1e1e24; --accent-blue: #00e5ff; --accent-purple: #b300ff; --text-main: #e0e0e0; }
        body { background-color: var(--bg-dark); color: var(--text-main); font-family: 'Segoe UI', Tahoma, sans-serif; margin: 0; display: flex; height: 100vh; overflow: hidden; }
        
        .sidebar { width: 280px; background-color: #0a0a0c; border-right: 1px solid #333; padding: 20px; display: flex; flex-direction: column; }
        .new-chat-btn { background: transparent; border: 1px solid var(--accent-blue); color: var(--accent-blue); padding: 10px; border-radius: 5px; cursor: pointer; font-weight: bold; width: 100%; margin-bottom: 20px;}
        .new-chat-btn:hover { background: rgba(0, 229, 255, 0.1); }
        .history-item { display: flex; justify-content: space-between; align-items: center; padding: 10px; border-bottom: 1px solid #222; font-size: 0.9em; transition: 0.2s;}
        .history-item:hover { background: #1a1a20; }
        .history-btn { background: none; border: none; color: #888; cursor: pointer; padding: 2px 5px; font-size: 1.1em;}
        .history-btn:hover { color: var(--accent-blue); }
        
        .main-content { flex: 1; display: flex; flex-direction: column; height: 100%; position: relative; }
        .chat-container { flex: 1; padding: 20px; max-width: 1000px; margin: 0 auto; width: 100%; overflow-y: auto; display: flex; flex-direction: column; gap: 20px; padding-bottom: 130px;}
        
        .user-bubble { background-color: #2a2a35; padding: 15px 20px; border-radius: 15px 15px 0 15px; align-self: flex-end; max-width: 80%; border: 1px solid #444; margin-left: auto; box-shadow: 0 4px 6px rgba(0,0,0,0.3);}
        .ai-bubble { background-color: var(--bg-panel); padding: 15px 20px; border-radius: 15px 15px 15px 0; align-self: flex-start; max-width: 90%; width: 100%; border: 1px solid #333; box-shadow: 0 4px 6px rgba(0,0,0,0.3);}
        
        .code-header { background: #111; padding: 8px 15px; display: flex; justify-content: space-between; align-items: center; border-radius: 8px 8px 0 0; font-family: monospace; font-size: 0.85em; border: 1px solid #333; border-bottom: none;}
        .code-header button { background: #333; border: 1px solid #444; color: white; cursor: pointer; padding: 4px 10px; border-radius: 5px; margin-left: 8px; transition: 0.2s;}
        .code-header button:hover { background: var(--accent-blue); color: black; border-color: var(--accent-blue);}
        .markdown-rendered pre { margin-top: 0; border-radius: 0 0 8px 8px; border: 1px solid #333;}
        
        /* GEMINI WEB STYLE INPUT BAR */
        .input-area { position: absolute; bottom: 0; width: 100%; padding: 20px; background: linear-gradient(0deg, #121212 85%, transparent); display: flex; justify-content: center; }
        .input-wrapper { display: flex; align-items: flex-end; gap: 12px; width: 100%; max-width: 900px; background: #1e1e24; padding: 12px 15px; border-radius: 24px; border: 1px solid #444; }
        .input-field { background: transparent; border: none; color: white; resize: none; flex-grow: 1; font-size: 1em; max-height: 150px; overflow-y: auto; padding: 5px; font-family: inherit; }
        .input-field:focus { outline: none; }
        
        .file-btn { background: #2a2a35; border: 1px solid #444; color: var(--accent-blue); width: 38px; height: 38px; border-radius: 50%; display: flex; align-items: center; justify-content: center; cursor: pointer; font-size: 1.4em; transition: 0.2s; flex-shrink: 0; }
        .file-btn:hover { background: rgba(0, 229, 255, 0.2); border-color: var(--accent-blue); }
        
        .btn-submit { background: linear-gradient(135deg, var(--accent-blue), var(--accent-purple)); color: white; border: none; width: 38px; height: 38px; border-radius: 50%; display: flex; align-items: center; justify-content: center; cursor: pointer; font-size: 1.2em; flex-shrink: 0; transition: transform 0.1s; }
        .btn-submit:active { transform: scale(0.95); }
    """),
    
    Script("""
        function renderMarkdown() {
            document.querySelectorAll('.markdown-raw').forEach(el => {
                let container = el.nextElementSibling;
                container.innerHTML = marked.parse(el.textContent);
                el.classList.remove('markdown-raw');
                
                container.querySelectorAll('pre').forEach(pre => {
                    let code = pre.querySelector('code');
                    let lang = code.className.replace('language-', '') || 'texte';
                    
                    let header = document.createElement('div');
                    header.className = 'code-header';
                    header.innerHTML = `
                        <span style="color:#00e5ff; font-weight:bold;">${lang.toUpperCase()}</span> 
                        <div>
                            <button onclick="copyCode(this)">📋 Copier</button> 
                            <button onclick="downloadCode(this, '${lang}')">⬇️ Télécharger ZIP</button>
                        </div>
                    `;
                    pre.parentNode.insertBefore(header, pre);
                    hljs.highlightElement(code);
                });
            });
            let chat = document.getElementById('chat-history');
            chat.scrollTop = chat.scrollHeight;
        }

        function copyCode(btn) {
            let code = btn.parentElement.parentElement.nextElementSibling.querySelector('code').innerText;
            navigator.clipboard.writeText(code);
            btn.innerText = "✅ Copié !";
            setTimeout(() => btn.innerText = "📋 Copier", 2000);
        }

        function downloadCode(btn, lang) {
            let code = btn.parentElement.parentElement.nextElementSibling.querySelector('code').innerText;
            let zip = new JSZip();
            let ext = lang === 'python' ? 'py' : lang === 'javascript' ? 'js' : lang === 'html' ? 'html' : 'txt';
            zip.file(`Aethas38_Code.${ext}`, code);
            zip.generateAsync({type:"blob"}).then(function(content) {
                let link = document.createElement('a');
                link.href = URL.createObjectURL(content);
                link.download = `Aethas38_Export_${lang}.zip`;
                link.click();
            });
        }
    """)
)

app, rt = fast_app(hdrs=theme_hdrs, secret_key="super_secret_aethas")

# ==========================================
# SECTION 2. CONNEXION AUX IA & CALCUL COÛTS
# ==========================================
local_client = OpenAI(base_url="http://192.168.1.20:11434/v1", api_key="ollama")
groq_client = Groq(api_key=os.getenv("GROQ_API_KEY"))
gemini_client = genai.Client(api_key=os.getenv("GEMINI_API_KEY"))

def interroger_gemini_cascade(prompt):
    modeles_a_tester = ["gemini-3.1-pro-preview", "gemini-3.1-pro", "gemini-2.5-pro", "gemini-3.7-flash", "gemini-3.6-flash"]
    for nom_modele in modeles_a_tester:
        try:
            ia_response = gemini_client.models.generate_content(
                model=nom_modele,
                contents=prompt,
                config={"system_instruction": "Réponds toujours dans la même langue que celle utilisée par l'utilisateur."}
            )
            return ia_response.text, nom_modele 
        except Exception:
            continue 
    raise Exception("Tous les modèles Gemini sont indisponibles actuellement.")

def calculer_cout_ia(fournisseur, prompt_tokens, completion_tokens):
    if fournisseur == "GROQ":
        input_cost_per_m, output_cost_per_m = 0.59, 0.79
    elif fournisseur == "GOOGLE GEMINI":
        input_cost_per_m, output_cost_per_m = 0.075, 0.30
    else:
        return 0.00
    return round(((prompt_tokens * input_cost_per_m) + (completion_tokens * output_cost_per_m)) / 1_000_000, 6)

# ==========================================
# SECTION 3. INTERFACE UTILISATEUR & HISTORIQUE
# ==========================================
def render_user_msg(text):
    return Div(P(text, style="margin:0; font-weight:bold; white-space: pre-wrap;"), cls="user-bubble")

def render_ai_msg(text, intention, nom_fournisseur, nom_modele, couleur):
    encart_info = Div(
        Span(f"🎯 {intention}", style="margin-right: 10px; color:#aaa; font-size:0.8em;"),
        Span(f"🏢 {nom_fournisseur}", style="margin-right: 10px; color:#aaa; font-size:0.8em;"),
        Span(f"🧠 {nom_modele}", style=f"margin-right: 10px; color:{couleur}; font-weight:bold; font-size:0.8em;"),
        style="border-bottom: 1px solid #333; padding-bottom: 10px; margin-bottom: 15px; display: flex; flex-wrap: wrap; gap: 5px;"
    )
    return Div(
        encart_info,
        Div(text, cls="markdown-raw", style="display:none;"),
        Div(cls="markdown-rendered", style="line-height: 1.6;"),
        cls="ai-bubble", style=f"border-left: 4px solid {couleur};"
    )

def render_sidebar():
    history = load_history()
    sorted_chats = sorted(history.items(), key=lambda x: (not x[1].get('pinned', False), x[0]))
    
    items = []
    for chat_id, data in sorted_chats:
        pin_icon = "📌" if data.get('pinned', False) else "📍"
        items.append(
            Div(
                Span(data['title'], hx_post=f"/load_chat/{chat_id}", hx_target="body", style="flex:1; cursor:pointer; white-space: nowrap; overflow: hidden; text-overflow: ellipsis;"),
                Button(pin_icon, cls="history-btn", hx_post=f"/pin/{chat_id}", hx_target="#sidebar-content", hx_swap="outerHTML"),
                Button("✏️", cls="history-btn", hx_post=f"/rename/{chat_id}", hx_prompt="Nouveau nom :", hx_target="#sidebar-content", hx_swap="outerHTML"),
                Button("🗑️", cls="history-btn", hx_post=f"/delete/{chat_id}", hx_target="#sidebar-content", hx_swap="outerHTML"),
                cls="history-item", id=f"chat-item-{chat_id}"
            )
        )
    return Div(*items, id="sidebar-content", style="flex:1; overflow-y:auto;", hx_swap_oob="true")

@rt('/')
def get(session):
    if 'current_chat_id' not in session:
        session['current_chat_id'] = str(uuid.uuid4())
        
    chat_id = session['current_chat_id']
    history = load_history()
    
    messages_ui = []
    if chat_id in history:
        for msg in history[chat_id].get('messages', []):
            if msg['role'] == 'user':
                messages_ui.append(render_user_msg(msg['content']))
            else:
                messages_ui.append(render_ai_msg(msg['content'], msg.get('intention', ''), msg.get('fournisseur', ''), msg.get('modele', ''), msg.get('couleur', '#333')))
    
    if messages_ui:
        messages_ui.append(Script("setTimeout(renderMarkdown, 100);"))

    return Title("Aethas38 - Mine_AI"), Body(
        Div(
            Div(H2("AETHAS 38", style="color:#00e5ff; margin:0; text-align:center;"), P("Tech Core", style="font-size:0.7em; color:#888; text-align:center;"), margin_bottom="20px"),
            Form(Button("+ Nouvelle discussion", cls="new-chat-btn"), hx_post="/new_chat", hx_target="body"),
            render_sidebar(),
            Div(P("👤 Admin (Xavier)", style="color: #aaa; margin: 0; text-align:center;"), style="border-top: 1px solid #333; padding-top: 15px;"),
            cls="sidebar"
        ),
        
        Div(
            Div(*messages_ui, id="chat-history", cls="chat-container"),
            Div(
                Form(
                    Div(
                        # Bouton + (Label déclenchant l'input file caché supportant le drag & drop et jusqu'à 10 fichiers)
                        Label("+", _for="file-upload", cls="file-btn", title="Importer des fichiers (max 10)"),
                        Input(type="file", name="files", id="file-upload", multiple=True, style="display:none;"),
                        # Textarea Gemini Web style avec ascenseur automatique
                        Textarea(name="prompt", id="prompt-input", placeholder="Entrez votre requête ou déposez vos fichiers ici...", cls="input-field", rows="1"),
                        # Bouton Flèche de validation
                        Button("↑", type="submit", cls="btn-submit", title="Exécuter"),
                        cls="input-wrapper"
                    ),
                    hx_post="/chat",
                    hx_target="#chat-history",
                    hx_swap="beforeend",
                    hx_on__after_request="this.reset()",
                    enctype="multipart/form-data",
                    style="width: 100%; display: flex; justify-content: center;"
                ),
                cls="input-area"
            ),
            cls="main-content"
        )
    )

# ==========================================
# SECTION 4. LE MOTEUR DE ROUTAGE & UPLOAD MULTI-FICHIERS
# ==========================================
@rt('/chat')
async def post(request, session: dict):
    try:
        form = await request.form()
        prompt = form.get("prompt", "")
        files = form.getlist("files") # Récupération de la liste des fichiers (jusqu'à 10)
        
        chat_id = session.get('current_chat_id')
        history = load_history()
        
        if chat_id not in history:
            history[chat_id] = {"title": (prompt[:30] if prompt else "Fichiers joints") + "...", "pinned": False, "messages": []}

        file_contents_summary = []
        file_names = []
        
        if files:
            for file in files[:10]: # Limitation stricte à 10 fichiers
                if file and file.filename:
                    file_names.append(file.filename)
                    content_bytes = await file.read()
                    try:
                        text_content = content_bytes.decode('utf-8', errors='ignore')
                    except Exception:
                        text_content = "[Fichier binaire ou non textuel]"
                    file_contents_summary.append(f"\n--- Fichier: {file.filename} ---\n{text_content}")

        full_prompt = prompt if prompt else f"Analyse ces fichiers : {', '.join(file_names)}"
        if file_contents_summary:
            full_prompt += "\n".join(file_contents_summary)

        # Analyse par l'Aiguilleur
        router_prompt = """Tu es le Routeur. Choisis : "groq" (code/maths), "gemini" (création/rédaction), "local" (salutations).
        Format strict : {"intention": "...", "fournisseur": "..."}. Ne dis rien d'autre."""
        router_response = local_client.chat.completions.create(model="qwen2.5:3b", messages=[{"role": "system", "content": router_prompt}, {"role": "user", "content": full_prompt}], temperature=0.1)
        
        raw_decision = router_response.choices[0].message.content
        start_idx = raw_decision.find('{')
        end_idx = raw_decision.rfind('}') + 1
        decision = json.loads(raw_decision[start_idx:end_idx]) if start_idx != -1 and end_idx != 0 else {"intention": "Secours", "fournisseur": "gemini"}
        fournisseur = decision.get("fournisseur", "local").lower()

        # Requête IA & Tokens
        prompt_tokens = len(full_prompt.split()) * 2
        completion_tokens = 0

        if fournisseur == "groq":
            ia_response = groq_client.chat.completions.create(messages=[{"role": "system", "content": "Réponds dans la langue de l'utilisateur."}, {"role": "user", "content": full_prompt}], model="llama-3.3-70b-versatile")
            reponse_finale, couleur, nom_fournisseur, nom_modele = ia_response.choices[0].message.content, "#f97316", "GROQ", "llama-3.3-70b"
            if hasattr(ia_response, 'usage') and ia_response.usage:
                prompt_tokens = ia_response.usage.prompt_tokens
                completion_tokens = ia_response.usage.completion_tokens
        elif fournisseur == "gemini":
            reponse_finale, nom_modele = interroger_gemini_cascade(full_prompt)
            couleur, nom_fournisseur = "#3b82f6", "GOOGLE GEMINI"
            completion_tokens = len(reponse_finale.split()) * 2
        else:
            ia_response = local_client.chat.completions.create(model="qwen2.5:3b", messages=[{"role": "system", "content": "Réponds dans la langue de l'utilisateur."}, {"role": "user", "content": full_prompt}])
            reponse_finale, couleur, nom_fournisseur, nom_modele = ia_response.choices[0].message.content, "#10b981", "SERVEUR LOCAL", "qwen2.5:3b"
            if hasattr(ia_response, 'usage') and ia_response.usage:
                prompt_tokens = ia_response.usage.prompt_tokens
                completion_tokens = ia_response.usage.completion_tokens

        # Sauvegarde en JSON
        display_msg = f"{prompt}" + (f"\n📎 [Fichiers: {', '.join(file_names)}]" if file_names else "")
        history[chat_id]['messages'].append({"role": "user", "content": display_msg})
        history[chat_id]['messages'].append({"role": "ai", "content": reponse_finale, "intention": decision.get('intention', ''), "fournisseur": nom_fournisseur, "modele": nom_modele, "couleur": couleur})
        save_history(history)

        user_bubble = render_user_msg(display_msg)
        ai_bubble = render_ai_msg(reponse_finale, decision.get('intention', ''), nom_fournisseur, nom_modele, couleur)

        return user_bubble, ai_bubble, Script("renderMarkdown();"), render_sidebar()

    except Exception as e:
        return Div(f"Erreur : {e}", style="color:red; background:#300; padding:15px; border-radius:10px;")

# ==========================================
# SECTION 5. ROUTES DE L'HISTORIQUE & NAVIGATION (RedirectResponse FIX)
# ==========================================
@rt('/new_chat')
def post(session):
    session['current_chat_id'] = str(uuid.uuid4())
    return RedirectResponse('/', status_code=303)

@rt('/load_chat/{chat_id}')
def post(chat_id: str, session):
    session['current_chat_id'] = chat_id
    return RedirectResponse('/', status_code=303)

@rt('/rename/{chat_id}')
def post(chat_id: str, req):
    nouveau_nom = req.headers.get("HX-Prompt")
    if nouveau_nom: 
        history = load_history()
        if chat_id in history:
            history[chat_id]['title'] = nouveau_nom
            save_history(history)
    return render_sidebar()

@rt('/pin/{chat_id}')
def post(chat_id: str):
    history = load_history()
    if chat_id in history:
        history[chat_id]['pinned'] = not history[chat_id].get('pinned', False)
        save_history(history)
    return render_sidebar()

@rt('/delete/{chat_id}')
def post(chat_id: str):
    history = load_history()
    if chat_id in history:
        del history[chat_id]
        save_history(history)
    return render_sidebar()

if __name__ == '__main__':
    serve(port=5001)