import requests
import json
import time
import os
import re
import sys
from typing import List, Dict, Any

"""
v3_multi.py
-- Conversație multi-AI pe mai multe mașini (Ollama) cu roluri și reguli de noutate --

Noutăți față de v2.py:
- Suport pentru N mașini (>=2), fiecare cu IP, model, rol, timeout și limită de tokeni.
- Roluri: generator / critic / moderator / sintetizator (synth). Promptare specifică fiecărui rol.
- Reguli de noutate: fiecare răspuns trebuie să adauge cel puțin o idee nouă față de registrul global.
- Rezumat de idei menționate este injectat în prompt ca să minimizeze repetițiile.
- Log detaliat per rundă și per mașină, cu număr de tokeni (dacă API-ul returnează eval_count).
- Retry + exponential backoff pe timeouts/erori.
- Întreținere registru de fraze/idei (deduplicate pe fraze curate).

Necesită: servere Ollama accesibile la /api/chat pe fiecare IP:PORT.
"""

PORT_DEFAULT = "11434"
MAX_RETRY = 3
LOG_FILE = os.path.join(os.path.dirname(__file__), "conversatie.txt")
REGISTRY_FILE = os.path.join(os.path.dirname(__file__), "idei.json")
INTREBARE_IMPLICITA = "Ce idei concrete, ne-repetate, poți adăuga mai departe?"

# Config implicită globală și opțiuni avansate Ollama
GLOBAL_DEFAULTS = {
    "runde": 5,                    # 0 = infinit
    "pauza": 3,                    # secunde între runde
    "context_messages_limit": 50,  # câte mesaje din istoric păstrăm
    "stream": True,                # stream răspunsuri
    # Opțiuni Ollama comune (pot fi suprascrise per mașină)
    "options": {
        "num_ctx": 4096,
        "temperature": 0.7,
        "top_p": 0.9,
        "repeat_penalty": 1.1
    }
}

# ---- Utilitare ----

def validare_ip(ip_str: str) -> bool:
    pattern = r'^(?:(?:25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)\.){3}(?:25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)$'
    return re.match(pattern, ip_str) is not None

def write_to_log(message: str) -> None:
    with open(LOG_FILE, "a", encoding="utf-8") as f:
        f.write(message + "\n")

def norm_text(s: str) -> str:
    s = s.lower()
    s = re.sub(r'\s+', ' ', s).strip()
    return s

def extrage_fraze_idei(text: str) -> List[str]:
    """
    Heuristic simplu: împarte pe punct/linie nouă/bullet și returnează fraze curate,
    ignorând linii foarte scurte (<5 caractere).
    """
    chunks = re.split(r'[.\n•\-\u2022;]+', text)
    out = []
    for c in chunks:
        t = norm_text(c)
        if len(t) >= 5:
            out.append(t)
    return out

def salveaza_registru(registry: Dict[str, bool]) -> None:
    try:
        with open(REGISTRY_FILE, "w", encoding="utf-8") as f:
            json.dump(sorted(list(registry.keys())), f, ensure_ascii=False, indent=2)
        write_to_log(f"Registru idei salvat în {REGISTRY_FILE}")
    except Exception as e:
        write_to_log(f"Eroare la salvarea registrului: {e}")

# ---- Promptare pe roluri ----

def sistem_role_prompt(subiect: str, rol: str, max_tokens: int) -> str:
    baza = (
        f"Discuți despre: {subiect}. Răspunde concis (max {max_tokens} tokeni), clar, în română.\n"
        "Reguli generale:\n"
        "1) Evită repetițiile. 2) Adaugă cel puțin O IDEE NOUĂ față de lista 'Idei deja menționate'. "
        "3) Listează ideile ca bullets scurte, concrete, testabile. 4) Fii specific (cifre, condiții, costuri, pași)."
    )
    if rol == "generator":
        return baza + "\nRol: GENERATOR – propune idei noi, combinații neobișnuite, dar plauzibile."
    if rol == "critic":
        return baza + "\nRol: CRITIC – semnalează riscuri/limitări și oferă o îmbunătățire alternativă pentru fiecare punct."
    if rol == "moderator":
        return baza + "\nRol: MODERATOR – impune noutatea, rezumă scurt progresul, pune o întrebare-țintă pentru runda următoare."
    if rol in ("synth", "sintetizator", "synthesizer"):
        return baza + "\nRol: SINTETIZATOR – combină cele mai bune idei, elimină redundanța, propune o mini-foaie de parcurs."
    # fallback
    return baza + "\nRol: GENERAL – adaugă idei utile și noi."

def construieste_mesaje(conversatie: List[Dict[str, str]], subiect: str, rol: str, max_tokens: int, idei_deja: List[str]) -> List[Dict[str, str]]:
    sistem = sistem_role_prompt(subiect, rol, max_tokens)
    idei_block = "Idei deja menționate până acum:\n- " + "\n- ".join(idei_deja[-30:]) if idei_deja else "Nu există idei înregistrate încă."
    instructiune = (
        f"{idei_block}\n"
        "Asigură-te că fiecare punct nou este ne-repetat față de listă. "
        f"După idei, încheie cu: 'NEXT: {INTREBARE_IMPLICITA}'"
    )
    # Pregătim istoricul conversației + instructiunea de rol
    messages = [{"role": "system", "content": sistem}]
    messages.extend(conversatie)
    messages.append({"role": "user", "content": instructiune})
    return messages

# ---- Rețea/Request ----

def api_chat_stream(url: str, payload: Dict[str, Any], timeout: int) -> Dict[str, Any]:
    retry = 0
    current_timeout = timeout
    raspuns_complet = ""
    tokens_eval = None
    while retry < MAX_RETRY:
        try:
            print(f"Încerc să trimit cerere la: http://{url}/api/chat cu modelul {payload.get('model')} (timeout: {current_timeout}s)")
            write_to_log(f"Încerc să trimit cerere la: http://{url}/api/chat cu modelul {payload.get('model')} (timeout: {current_timeout}s)")
            # Respectă opțiunea stream din payload
            use_stream = bool(payload.get("stream", True))
            resp = requests.post(f"http://{url}/api/chat", json=payload, timeout=current_timeout, stream=use_stream)
            resp.raise_for_status()
            if use_stream:
                for line in resp.iter_lines():
                    if not line:
                        continue
                    chunk = json.loads(line)
                    if "message" in chunk and "content" in chunk["message"]:
                        token = chunk["message"]["content"]
                        raspuns_complet += token
                        sys.stdout.write(token)
                        sys.stdout.flush()
                    if chunk.get("done", False):
                        tokens_eval = chunk.get("eval_count", None)
                        print()
                        break
                return {"text": raspuns_complet, "tokens": tokens_eval}
            else:
                # non-stream: citim JSON-ul complet
                chunk = resp.json()
                if "message" in chunk and "content" in chunk["message"]:
                    raspuns_complet = chunk["message"]["content"]
                tokens_eval = chunk.get("eval_count", None)
                print(raspuns_complet)
                return {"text": raspuns_complet, "tokens": tokens_eval}
        except requests.exceptions.ReadTimeout:
            msg = f"Timeout depășit ({current_timeout}s) către {url}. Reîncerc..."
            print(msg); write_to_log(msg)
            retry += 1; current_timeout *= 2
        except requests.exceptions.RequestException as e:
            msg = f"Eroare către {url}: {e}"
            print(msg); write_to_log(msg)
            retry += 1; current_timeout *= 2
    msg = f"Maxim de încercări atins ({MAX_RETRY}) către {url}. Abort."
    print(msg); write_to_log(msg)
    return {"text": None, "tokens": None}

def test_server(url: str, timeout: int) -> bool:
    """Verifică dacă serverul Ollama răspunde pe /api/tags."""
    try:
        r = requests.get(f"http://{url}/api/tags", timeout=timeout)
        if r.status_code == 200:
            return True
        write_to_log(f"Server {url} răspunde cu cod {r.status_code}")
        return False
    except requests.exceptions.RequestException as e:
        write_to_log(f"Nu se poate atinge {url}: {e}")
        return False

# ---- Main ----

def main():
    # Inițializare log
    with open(LOG_FILE, "w", encoding="utf-8") as f:
        f.write("== Conversație multi-AI ==\n")

    subiect = input("Introdu tema de discuție: ").strip()

    # Alege modul de configurare: simplu sau avansat
    mod = input("Mod configurare [simplu/avansat] (implicit: simplu): ").strip().lower() or "simplu"

    if mod == "avansat":
        try:
            runde_input = input(f"Număr runde (0 = infinit, implicit {GLOBAL_DEFAULTS['runde']}): ").strip()
            numar_runde = int(runde_input) if runde_input else GLOBAL_DEFAULTS["runde"]
        except ValueError:
            print(f"Valoare invalidă. Implicit {GLOBAL_DEFAULTS['runde']} runde."); numar_runde = GLOBAL_DEFAULTS["runde"]

        try:
            pauza_input = input(f"Pauză între runde sec (implicit {GLOBAL_DEFAULTS['pauza']}): ").strip()
            PAUZA_INTRE_RUNDE_SEC = int(pauza_input) if pauza_input else GLOBAL_DEFAULTS["pauza"]
        except ValueError:
            print(f"Valoare invalidă. Implicit {GLOBAL_DEFAULTS['pauza']} secunde."); PAUZA_INTRE_RUNDE_SEC = GLOBAL_DEFAULTS["pauza"]

        try:
            ctx_input = input(f"Limită mesaje context (implicit {GLOBAL_DEFAULTS['context_messages_limit']}): ").strip()
            CONTEXT_LIMIT = int(ctx_input) if ctx_input else GLOBAL_DEFAULTS["context_messages_limit"]
        except ValueError:
            print(f"Valoare invalidă. Implicit {GLOBAL_DEFAULTS['context_messages_limit']}."); CONTEXT_LIMIT = GLOBAL_DEFAULTS["context_messages_limit"]

        # Opțiuni Ollama comune
        stream_raw = input(f"Stream răspunsuri [da/nu] (implicit {'da' if GLOBAL_DEFAULTS['stream'] else 'nu'}): ").strip().lower()
        STREAM_MODE = GLOBAL_DEFAULTS["stream"] if stream_raw not in ("da", "nu") else (stream_raw == "da")
        try:
            temp_in = input(f"Temperature (implicit {GLOBAL_DEFAULTS['options']['temperature']}): ").strip()
            TEMP = float(temp_in) if temp_in else GLOBAL_DEFAULTS["options"]["temperature"]
        except ValueError:
            TEMP = GLOBAL_DEFAULTS["options"]["temperature"]
        try:
            top_p_in = input(f"Top_p (implicit {GLOBAL_DEFAULTS['options']['top_p']}): ").strip()
            TOP_P = float(top_p_in) if top_p_in else GLOBAL_DEFAULTS["options"]["top_p"]
        except ValueError:
            TOP_P = GLOBAL_DEFAULTS["options"]["top_p"]
        try:
            rp_in = input(f"Repeat penalty (implicit {GLOBAL_DEFAULTS['options']['repeat_penalty']}): ").strip()
            REPEAT_PENALTY = float(rp_in) if rp_in else GLOBAL_DEFAULTS["options"]["repeat_penalty"]
        except ValueError:
            REPEAT_PENALTY = GLOBAL_DEFAULTS["options"]["repeat_penalty"]
        try:
            num_ctx_in = input(f"num_ctx (implicit {GLOBAL_DEFAULTS['options']['num_ctx']}): ").strip()
            NUM_CTX = int(num_ctx_in) if num_ctx_in else GLOBAL_DEFAULTS["options"]["num_ctx"]
        except ValueError:
            NUM_CTX = GLOBAL_DEFAULTS["options"]["num_ctx"]
    else:
        # modul simplu: folosește implicitele
        numar_runde = GLOBAL_DEFAULTS["runde"]
        PAUZA_INTRE_RUNDE_SEC = GLOBAL_DEFAULTS["pauza"]
        CONTEXT_LIMIT = GLOBAL_DEFAULTS["context_messages_limit"]
        STREAM_MODE = GLOBAL_DEFAULTS["stream"]
        TEMP = GLOBAL_DEFAULTS["options"]["temperature"]
        TOP_P = GLOBAL_DEFAULTS["options"]["top_p"]
        REPEAT_PENALTY = GLOBAL_DEFAULTS["options"]["repeat_penalty"]
        NUM_CTX = GLOBAL_DEFAULTS["options"]["num_ctx"]

    try:
        numar_masini = int(input("Câte mașini (>=2) dorești să configurezi?: "))
    except ValueError:
        print("Valoare invalidă. Implicit 2 mașini."); numar_masini = 2
    if numar_masini < 2:
        numar_masini = 2

    masini: List[Dict[str, Any]] = []
    for idx in range(numar_masini):
        print(f"\n-- Config mașină #{idx+1} --")
        if idx == 0:
            # Default-uri pentru prima mașină
            ip_default = "100.76.154.75"
            model_default = "qwen3-coder:480b-cloud"
            max_tokens_default = 700
            timeout_default = 120
        elif idx == 1:
            # Default-uri pentru a doua mașină
            ip_default = "100.73.213.29"
            model_default = "gpt-oss:120b-cloud"
            max_tokens_default = 700
            timeout_default = 120
        else:
            # Pentru mașini suplimentare, folosește default-urile originale
            ip_default = None
            model_default = "llama3:latest"
            max_tokens_default = 200
            timeout_default = 120

        while True:
            ip_prompt = f"IP (default {ip_default}): " if ip_default else "IP (ex: 100.76.154.75): "
            ip_raw = input(ip_prompt).strip()
            if not ip_raw and ip_default:
                ip_raw = ip_default
            if validare_ip(ip_raw):
                break
            print("IP invalid, reîncearcă.")
        
        port = input(f"Port (gol pentru default {PORT_DEFAULT}): ").strip() or PORT_DEFAULT
        model_input = input(f"Model (default {model_default}): ").strip()
        model = model_default if not model_input else model_input
        rol = input("Rol [generator/critic/moderator/synth] (implicit: generator): ").strip().lower() or "generator"
        try:
            max_tokens_input = input(f"Limita tokeni răspuns (default {max_tokens_default}): ") or str(max_tokens_default)
            max_tokens = int(max_tokens_input)
        except ValueError:
            max_tokens = max_tokens_default
        try:
            timeout_input = input(f"Timeout inițial sec (default {timeout_default}): ") or str(timeout_default)
            timeout = int(timeout_input)
        except ValueError:
            timeout = timeout_default

        # Opțiuni avansate per mașină (dacă modul avansat)
        if mod == "avansat":
            try:
                temp_in_m = input(f"Temperature pentru {model} (gol = {TEMP}): ").strip()
                temp_m = float(temp_in_m) if temp_in_m else TEMP
            except ValueError:
                temp_m = TEMP
            try:
                top_p_in_m = input(f"Top_p pentru {model} (gol = {TOP_P}): ").strip()
                top_p_m = float(top_p_in_m) if top_p_in_m else TOP_P
            except ValueError:
                top_p_m = TOP_P
            try:
                rp_in_m = input(f"Repeat penalty pentru {model} (gol = {REPEAT_PENALTY}): ").strip()
                rp_m = float(rp_in_m) if rp_in_m else REPEAT_PENALTY
            except ValueError:
                rp_m = REPEAT_PENALTY
            try:
                num_ctx_in_m = input(f"num_ctx pentru {model} (gol = {NUM_CTX}): ").strip()
                num_ctx_m = int(num_ctx_in_m) if num_ctx_in_m else NUM_CTX
            except ValueError:
                num_ctx_m = NUM_CTX
        else:
            temp_m = TEMP; top_p_m = TOP_P; rp_m = REPEAT_PENALTY; num_ctx_m = NUM_CTX

        masini.append({
            "url": f"{ip_raw}:{port}",
            "model": model,
            "rol": rol,
            "max_tokens": max_tokens,
            "timeout": timeout,
            "options": {
                "temperature": temp_m,
                "top_p": top_p_m,
                "repeat_penalty": rp_m,
                "num_ctx": num_ctx_m
            }
        })

    # Istoric conversație comun
    conversatie: List[Dict[str, str]] = []
    # Registru idei unice (fraze normalizate)
    registry: Dict[str, bool] = {}

    # Test conexiune pentru fiecare server
    for m in masini:
        if not test_server(m["url"], m["timeout"]):
            print(f"Avertisment: Serverul {m['url']} nu răspunde la /api/tags. Continuăm oricum.")

    runda = 0
    while True:
        runda += 1
        if numar_runde != 0 and runda > numar_runde:
            done = f"Procesul s-a încheiat după {numar_runde} runde."
            print(done); write_to_log(done)
            salveaza_registru(registry)
            break
        print(f"\n===== RUNDA {runda} ====="); write_to_log(f"\n===== RUNDA {runda} =====")

        # parcurgere round-robin
        for m in masini:
            rol = m["rol"]
            max_tok = m["max_tokens"]
            timeout = m["timeout"]
            url = m["url"]
            model = m["model"]

            idei_list = list(registry.keys())
            # Trunchiere istoric conversație la ultimul CONTEXT_LIMIT
            conv_trunchiat = conversatie[-CONTEXT_LIMIT:] if len(conversatie) > CONTEXT_LIMIT else conversatie
            mesaje = construieste_mesaje(conv_trunchiat, subiect, rol, max_tok, idei_list)
            payload = {
                "model": model,
                "messages": mesaje,
                "stream": STREAM_MODE,
                "options": {
                    "num_predict": max_tok,
                    "temperature": m["options"]["temperature"],
                    "top_p": m["options"]["top_p"],
                    "repeat_penalty": m["options"]["repeat_penalty"],
                    "num_ctx": m["options"]["num_ctx"]
                }
            }

            rezultat = api_chat_stream(url, payload, timeout)
            text = rezultat["text"]
            tokens = rezultat["tokens"]

            if text is None:
                stop = f"Nu s-a primit răspuns de la {url}. Oprire."
                print(stop); write_to_log(stop)
                salveaza_registru(registry)
                return

            write_to_log(f"[{rol.upper()} @ {url} :: tokens={tokens}] {text}")

            # actualizează registry de idei
            fraze = extrage_fraze_idei(text)
            adaugate = 0
            for fr in fraze:
                if fr not in registry:
                    registry[fr] = True
                    adaugate += 1

            # adaugă în conversație ca mesaj al 'assistant'-ului și întrebare implicită
            conversatie.append({"role": "assistant", "content": text})
            conversatie.append({"role": "user", "content": INTREBARE_IMPLICITA})

            print(f"\n[INFO] Idei noi adăugate în registru: {adaugate}")
            write_to_log(f"[INFO] Idei noi adăugate: {adaugate}")

        # Mic meniu între runde (ajustează parametri globali rapid)
        cmd = input(f"\nEnter pentru continuare, sau 'setari' pentru ajustări: ").strip().lower()
        if cmd == "setari":
            try:
                p = input(f"Pauză sec (actual {PAUZA_INTRE_RUNDE_SEC}): ").strip()
                if p:
                    PAUZA_INTRE_RUNDE_SEC = int(p)
            except ValueError:
                print("Valoare pauză invalidă, păstrăm setarea curentă.")
            try:
                cl = input(f"Limită mesaje context (actual {CONTEXT_LIMIT}): ").strip()
                if cl:
                    CONTEXT_LIMIT = int(cl)
            except ValueError:
                print("Valoare limită context invalidă.")
            sm = input(f"Stream [da/nu] (actual {'da' if STREAM_MODE else 'nu'}): ").strip().lower()
            if sm in ("da", "nu"):
                STREAM_MODE = (sm == "da")
        time.sleep(PAUZA_INTRE_RUNDE_SEC)

if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        print("\nOprire prin Ctrl+C.")