import httpx
import json
import os
import sys
import time
import argparse

# Configuration
API_URL = "https://localhost:7860/v1/chat/completions"
SESSION_ID = "aab82b16"
SESSION_FILE = "/run/media/julian/ML4/ollama-work/all_space/sessions/aab82b16.json"

def load_local_session():
    if not os.path.exists(SESSION_FILE):
        return {"session_id": SESSION_ID, "history": []}
    with open(SESSION_FILE, "r") as f:
        return json.load(f)

def save_local_session(history):
    data = load_local_session()
    data["history"] = history
    data["updated_at"] = str(time.time())
    with open(SESSION_FILE, "w") as f:
        json.dump(data, f, indent=2)

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--message", type=str, help="The message to send to the model")
    parser.add_argument("--preset", type=str, default="SUBJECTIVE", help="The PX preset (e.g. SUBJECTIVE, RESONANCE_CITY)")
    args = parser.parse_args()

    print("="*60)
    print(f" LIVE SPACE INTERFACE - SESSION: {SESSION_ID} ")
    print(f" MODE: {args.preset} ")
    print("="*60)
    
    session_data = load_local_session()
    history = session_data.get("history", [])
    model_id = session_data.get("model_id", "gemma3-1b-it")
    
    # Show last context
    if history:
        last_msg = history[-1]
        print(f"\n[LETZTER KONTEXT - {last_msg['role'].upper()}]:")
        content = last_msg['content']
        if isinstance(content, list):
             text = "".join([b.get("text", "") for b in content if b.get("type") == "text"])
             print(text[:300] + "..." if len(text) > 300 else text)
        else:
             print(content[:300] + "..." if len(content) > 300 else content)

    if args.message:
        new_user_msg = args.message
    else:
        new_user_msg = "Bitte setze deine Gedanken fort."

    print("\n" + "-"*20 + " MEIN INPUT (GEMINI CLI) " + "-"*20)
    print(new_user_msg)
    print("-" * 60)
    print("\n[WARTE AUF ANTWORT VON ALL_SPACE...]\n")

    # Prepare API payload
    api_messages = []
    for msg in history:
        role = msg["role"]
        content = msg["content"]
        if isinstance(content, list):
            text = "".join([b.get("text", "") for b in content if b.get("type") == "text"])
            content = text
        api_messages.append({"role": role, "content": content})
    
    api_messages.append({"role": "user", "content": new_user_msg})
    
    payload = {
        "model": model_id,
        "messages": api_messages,
        "px_subjective": True,
        "px_config_preset": args.preset,
        "temperature": 0.7,
        "max_tokens": 1024,
        "stream": True
    }

    full_response = ""
    print("[MODELL ANTWORT]:")
    try:
        with httpx.stream("POST", API_URL, json=payload, verify=False, timeout=None) as response:
            if response.status_code != 200:
                print(f"Error: {response.status_code}")
                print(response.read().decode())
                return

            for line in response.iter_lines():
                if not line:
                    continue
                if line.startswith("data: "):
                    data_str = line[6:]
                    if data_str == "[DONE]":
                        break
                    try:
                        data = json.loads(data_str)
                        delta = data["choices"][0]["delta"]
                        if "content" in delta:
                            content = delta["content"]
                            full_response += content
                            sys.stdout.write(content)
                            sys.stdout.flush()
                    except:
                        pass
    except Exception as e:
        print(f"\n[FEHLER]: {e}")
        return

    print("\n\n" + "="*60)
    
    # Save session
    new_history = history + [
        {"role": "user", "content": new_user_msg},
        {"role": "assistant", "content": full_response}
    ]
    save_local_session(new_history)
    print(f"Update: Session {SESSION_ID} gespeichert.")

if __name__ == "__main__":
    main()