#!/usr/bin/env python3
"""
Chat interactif en ligne de commande avec gemma4:26b via Ollama.

Usage interactif :
    python tests/test_image_chat_cli.py
    # puis taper des questions sur l'image fournie

Usage one-shot :
    python tests/test_image_chat_cli.py /chemin/vers/image.png "Que vois-tu ?"

Usage avec modèle différent :
    python tests/test_image_chat_cli.py --model qwen3-vl:8b image.png

Le script utilise l'API Ollama directement (via la lib `ollama` du projet,
`ollama==0.6.1` dans requirements.txt).
"""

import argparse
import base64
import sys
from pathlib import Path

try:
    import ollama
except ImportError:
    print("ERREUR : la librairie 'ollama' n'est pas installée.")
    print("Installez-la avec : pip install ollama")
    sys.exit(1)


DEFAULT_MODEL = "gemma4:26b"


def encode_image(image_path: str) -> str:
    """Encode une image en base64 pour l'API Ollama."""
    path = Path(image_path)
    if not path.exists():
        print(f"ERREUR : le fichier '{image_path}' n'existe pas.")
        sys.exit(1)
    if not path.is_file():
        print(f"ERREUR : '{image_path}' n'est pas un fichier.")
        sys.exit(1)
    with open(path, "rb") as f:
        return base64.b64encode(f.read()).decode("utf-8")


def get_client(host: str):
    """Renvoie un client Ollama configuré pour l'hôte donné."""
    return ollama.Client(host=host)


def check_ollama_running(host: str = "http://localhost:11434") -> bool:
    """Vérifie que le serveur Ollama est accessible."""
    try:
        client = get_client(host)
        client.list()
        return True
    except Exception as e:
        print(f"ERREUR : impossible de joindre Ollama sur {host}")
        print(f"Détail : {e}")
        print()
        print("Assurez-vous qu'Ollama est lancé :")
        print("  ollama serve")
        return False


def check_model_available(model: str, host: str = "http://localhost:11434") -> bool:
    """Vérifie que le modèle est disponible dans Ollama."""
    try:
        client = get_client(host)
        tags = client.list()
        # ollama.list() retourne un ListResponse avec un attribut 'models'
        models = getattr(tags, "models", [])

        model_names = []
        for m in models:
            if isinstance(m, dict):
                model_names.append(m.get("name", ""))
            else:
                model_names.append(getattr(m, "name", str(m)))

        # Correspondance exacte ou préfixe
        matched = [name for name in model_names if model in name]
        if matched:
            return True
        else:
            print(f"AVERTISSEMENT : modèle '{model}' non trouvé dans Ollama.")
            print(f"Modèles disponibles : {', '.join(model_names) or '(aucun)'}")
            print()
            print(f"Pour le télécharger :")
            print(f"  ollama pull {model}")
            return False
    except Exception as e:
        print(f"ERREUR : impossible de lister les modèles : {e}")
        return False


def chat_with_image(image_path: str, model: str, host: str = "http://localhost:11434") -> None:
    """Mode interactif : charge l'image une fois, puis pose des questions."""
    client = get_client(host)
    image_b64 = encode_image(image_path)
    print(f"🖼️  Image chargée : {image_path}")
    print(f"🤖 Modèle       : {model}")
    print(f"🔗 Ollama       : {host}")
    print()
    print("Mode interactif — tapez vos questions (ou 'exit'/'quit' pour sortir)")
    print("Tapez '/image /chemin/nouvelle.png' pour changer d'image")
    print("-" * 60)

    # Historique de conversation (sans l'image à chaque fois pour économiser la mémoire)
    messages = []

    while True:
        try:
            question = input("\nVous > ").strip()
        except (EOFError, KeyboardInterrupt):
            print("\n👋 Au revoir !")
            break

        if not question:
            continue

        if question.lower() in ("exit", "quit", "q"):
            print("👋 Au revoir !")
            break

        # Changement d'image
        if question.startswith("/image "):
            new_path = question[len("/image "):].strip()
            try:
                image_b64 = encode_image(new_path)
                image_path = new_path
                # Réinitialiser l'historique car image différente
                messages = []
                print(f"🖼️  Nouvelle image : {new_path}")
            except SystemExit:
                pass
            continue

        # Construire le message user avec l'image au premier tour
        # Ensuite, l'image n'est ré-envoyée que si l'historique est vide
        has_image_in_context = any(
            isinstance(m.get("images"), list) and len(m["images"]) > 0
            for m in messages
        )

        user_msg = {"role": "user", "content": question}
        if not has_image_in_context:
            # Première question ou image changée — inclure l'image
            user_msg["images"] = [image_b64]
        messages.append(user_msg)

        print(f"🤖 Réponse ({model})...", end=" ", flush=True)

        try:
            response = client.chat(
                model=model,
                messages=messages,
                stream=True,
                options={
                    "temperature": 0.2,
                    "num_predict": 2048,
                },
            )

            full_response = ""
            print()  # nouvelle ligne après le "..."
            for chunk in response:
                content = chunk.get("message", {}).get("content", "")
                if content:
                    print(content, end="", flush=True)
                    full_response += content

            print()  # retour à la ligne après la réponse
            messages.append({"role": "assistant", "content": full_response})

        except Exception as e:
            print(f"\n❌ Erreur : {e}")
            # Retirer le dernier message user en cas d'erreur
            messages.pop()


def one_shot(image_path: str, question: str, model: str, host: str = "http://localhost:11434") -> None:
    """Mode one-shot : une question, une réponse."""
    client = get_client(host)
    image_b64 = encode_image(image_path)

    messages = [
        {"role": "user", "content": question, "images": [image_b64]},
    ]

    try:
        response = client.chat(
            model=model,
            messages=messages,
            stream=True,
            options={
                "temperature": 0.2,
                "num_predict": 2048,
            },
        )

        print(f"🤖 {model} — '{question}'\n")
        for chunk in response:
            content = chunk.get("message", {}).get("content", "")
            if content:
                print(content, end="", flush=True)
        print()

    except Exception as e:
        print(f"❌ Erreur : {e}")
        sys.exit(1)


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Chat interactif avec une image via Ollama (gemma4:26b par défaut)",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Exemples :
  # Mode interactif avec une image
  python tests/test_image_chat_cli.py screenshot.png

  # Mode one-shot (question directe)
  python tests/test_image_chat_cli.py screenshot.png "Quels boutons vois-tu ?"

  # Avec un autre modèle
  python tests/test_image_chat_cli.py --model qwen3-vl:8b screenshot.png

  # Ollama sur une machine distante
  python tests/test_image_chat_cli.py --host http://dgx:11434 screenshot.png
        """,
    )
    parser.add_argument(
        "image",
        nargs="?",
        help="Chemin vers l'image à analyser",
    )
    parser.add_argument(
        "question",
        nargs="?",
        default=None,
        help="Question one-shot (si absent → mode interactif)",
    )
    parser.add_argument(
        "--model",
        default=DEFAULT_MODEL,
        help=f"Modèle Ollama à utiliser (défaut: {DEFAULT_MODEL})",
    )
    parser.add_argument(
        "--host",
        default="http://localhost:11434",
        help="URL du serveur Ollama (défaut: http://localhost:11434)",
    )

    args = parser.parse_args()

    # Vérifications préalables
    if not check_ollama_running(args.host):
        sys.exit(1)

    if not check_model_available(args.model, args.host):
        sys.exit(1)

    if not args.image:
        print("Utilisation interactive — veuillez fournir le chemin d'une image.")
        print()
        print("Usage :")
        print(f"  python {sys.argv[0]} /chemin/vers/image.png")
        print(f"  python {sys.argv[0]} /chemin/vers/image.png \"Votre question\"")
        print()
        parser.print_help()
        sys.exit(1)

    if args.question:
        # Mode one-shot
        one_shot(args.image, args.question, args.model, args.host)
    else:
        # Mode interactif
        chat_with_image(args.image, args.model, args.host)


if __name__ == "__main__":
    main()