from flask import Flask, Response, request
from flask import stream_with_context
import json
import time

# Refolosim logica din v4.py
import v4 as core

app = Flask(__name__)


def sse_event(event: dict) -> str:
    return f"data: {json.dumps(event, ensure_ascii=False)}\n\n"


def normalize_machines(machines):
    normalized = []
    for m in machines:
        ip = m.get("ip", "")
        port = str(m.get("port", core.PORT_DEFAULT))
        url = f"{ip}:{port}"
        normalized.append({
            "url": url,
            "model": m.get("model", "llama3:latest"),
            "rol": (m.get("rol", "generator") or "generator").lower(),
            "max_tokens": int(m.get("max_tokens", 300)),
            "timeout": int(m.get("timeout", 120)),
            "options": {
                "temperature": float(m.get("temperature", 0.7)),
                "top_p": float(m.get("top_p", 0.9)),
                "repeat_penalty": float(m.get("repeat_penalty", 1.1)),
                "num_ctx": int(m.get("num_ctx", 4096)),
            },
        })
    return normalized


def run_conversation(config: dict):
    subiect = config.get("subiect", "")
    numar_runde = int(config.get("runde", 5))
    pauza = int(config.get("pauza", 3))
    context_limit = int(config.get("context_limit", 50))
    stream_mode = bool(config.get("stream", False))  # folosim non-stream pentru răspunsuri compacte
    machines = normalize_machines(config.get("masini", []))

    conversatie = []
    registry = {}

    # Test servere
    for m in machines:
        ok = core.test_server(m["url"], m["timeout"])
        yield sse_event({"type": "server_check", "url": m["url"], "ok": ok})

    runda = 0
    while True:
        runda += 1
        if numar_runde != 0 and runda > numar_runde:
            yield sse_event({"type": "done", "message": f"Procesul s-a încheiat după {numar_runde} runde."})
            core.salveaza_registru(registry)
            break

        yield sse_event({"type": "round_start", "round": runda})

        for m in machines:
            rol = m["rol"]
            max_tok = m["max_tokens"]
            timeout = m["timeout"]
            url = m["url"]
            model = m["model"]

            idei_list = list(registry.keys())
            conv_trunchiat = conversatie[-context_limit:] if len(conversatie) > context_limit else conversatie
            mesaje = core.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 = core.api_chat_stream(url, payload, timeout)
            text = rezultat["text"]
            tokens = rezultat["tokens"]

            if text is None:
                yield sse_event({"type": "error", "url": url, "message": "Nu s-a primit răspuns."})
                core.salveaza_registru(registry)
                return

            fraze = core.extrage_fraze_idei(text)
            adaugate = 0
            for fr in fraze:
                if fr not in registry:
                    registry[fr] = True
                    adaugate += 1

            conversatie.append({"role": "assistant", "content": text})
            conversatie.append({"role": "user", "content": core.INTREBARE_IMPLICITA})

            yield sse_event({
                "type": "machine_response",
                "round": runda,
                "url": url,
                "rol": rol,
                "model": model,
                "tokens": tokens,
                "text": text,
                "new_ideas": adaugate,
            })

        yield sse_event({"type": "round_end", "round": runda, "pause_sec": pauza})
        time.sleep(pauza)


@app.get("/")
def index():
    # Servim o pagină simplă cu setări
    html = """
<!doctype html>
<html lang="ro">
<head>
  <meta charset="utf-8" />
  <meta name="viewport" content="width=device-width, initial-scale=1" />
  <title>Discuție multi-AI (browser)</title>
  <style>
    body { font-family: system-ui, Arial, sans-serif; margin: 20px; }
    h1 { margin-bottom: 0.5rem; }
    fieldset { border: 1px solid #ccc; padding: 10px; margin-bottom: 10px; }
    label { display: inline-block; min-width: 160px; margin: 4px 0; }
    input, select { padding: 4px 6px; }
    .row { display: flex; gap: 10px; flex-wrap: wrap; }
    .machines { margin-top: 10px; }
    .machine { border: 1px solid #ddd; padding: 8px; margin-bottom: 8px; }
    .log { white-space: pre-wrap; background: #fafafa; border: 1px solid #eee; padding: 10px; height: 300px; overflow: auto; }
    .controls { margin: 10px 0; }
    .ok { color: #0a7b00; }
    .err { color: #b00020; }
  </style>
</head>
<body>
  <h1>Discuție multi-AI (Ollama) – UI web</h1>
  <fieldset>
    <legend>Setări generale</legend>
    <div class="row">
      <label>Subiect: <input id="subiect" placeholder="Tema de discuție" /></label>
      <label>Runde (0=inf): <input id="runde" type="number" value="5" /></label>
      <label>Pauză (sec): <input id="pauza" type="number" value="3" /></label>
      <label>Limită context: <input id="context_limit" type="number" value="50" /></label>
      <label>Stream: 
        <select id="stream">
          <option value="false" selected>Nu</option>
          <option value="true">Da</option>
        </select>
      </label>
    </div>
  </fieldset>

  <fieldset class="machines">
    <legend>Mașini</legend>
    <div id="machines"></div>
    <div class="controls">
      <button id="addMachine">Adaugă mașină</button>
    </div>
  </fieldset>

  <div class="controls">
    <button id="start">Start conversație</button>
    <button id="stop" disabled>Stop</button>
  </div>

  <h3>Log</h3>
  <div id="log" class="log"></div>

  <script>
    const machinesDiv = document.getElementById('machines');
    const logDiv = document.getElementById('log');
    let es = null;

    function addMachine(defaults={}) {
      const wrapper = document.createElement('div');
      wrapper.className = 'machine';
      wrapper.innerHTML = `
        <div class="row">
          <label>IP: <input class="ip" placeholder="100.76.154.75" value="${defaults.ip||''}" /></label>
          <label>Port: <input class="port" type="number" value="${defaults.port||11434}" /></label>
          <label>Model: <input class="model" value="${defaults.model||'llama3:latest'}" /></label>
          <label>Rol: 
            <select class="rol">
              <option value="generator" selected>generator</option>
              <option value="critic">critic</option>
              <option value="moderator">moderator</option>
              <option value="synth">synth</option>
            </select>
          </label>
          <label>num_predict: <input class="max_tokens" type="number" value="${defaults.max_tokens||300}" /></label>
          <label>Timeout: <input class="timeout" type="number" value="${defaults.timeout||120}" /></label>
        </div>
        <div class="row">
          <label>temperature: <input class="temperature" type="number" step="0.1" value="${defaults.temperature||0.7}" /></label>
          <label>top_p: <input class="top_p" type="number" step="0.05" value="${defaults.top_p||0.9}" /></label>
          <label>repeat_penalty: <input class="repeat_penalty" type="number" step="0.1" value="${defaults.repeat_penalty||1.1}" /></label>
          <label>num_ctx: <input class="num_ctx" type="number" value="${defaults.num_ctx||4096}" /></label>
        </div>
        <div class="controls"><button class="remove">Șterge</button></div>
      `;
      wrapper.querySelector('.remove').onclick = () => wrapper.remove();
      machinesDiv.appendChild(wrapper);
    }

    document.getElementById('addMachine').onclick = () => addMachine();
    // două mașini implicite
    addMachine({ip:'100.76.154.75', model:'qwen3-coder:480b-cloud', max_tokens:700, timeout:120});
    addMachine({ip:'100.73.213.29', model:'gpt-oss:120b-cloud', max_tokens:700, timeout:120});

    function collectConfig() {
      const ms = Array.from(machinesDiv.querySelectorAll('.machine')).map(m => ({
        ip: m.querySelector('.ip').value.trim(),
        port: m.querySelector('.port').value.trim(),
        model: m.querySelector('.model').value.trim(),
        rol: m.querySelector('.rol').value,
        max_tokens: Number(m.querySelector('.max_tokens').value),
        timeout: Number(m.querySelector('.timeout').value),
        temperature: Number(m.querySelector('.temperature').value),
        top_p: Number(m.querySelector('.top_p').value),
        repeat_penalty: Number(m.querySelector('.repeat_penalty').value),
        num_ctx: Number(m.querySelector('.num_ctx').value),
      }));
      return {
        subiect: document.getElementById('subiect').value.trim(),
        runde: Number(document.getElementById('runde').value),
        pauza: Number(document.getElementById('pauza').value),
        context_limit: Number(document.getElementById('context_limit').value),
        stream: document.getElementById('stream').value === 'true',
        masini: ms,
      };
    }

    function log(msg, cls='') {
      const line = document.createElement('div');
      if (cls) line.className = cls;
      line.textContent = msg;
      logDiv.appendChild(line);
      logDiv.scrollTop = logDiv.scrollHeight;
    }

    document.getElementById('start').onclick = () => {
      const cfg = collectConfig();
      const qs = encodeURIComponent(JSON.stringify(cfg));
      if (es) es.close();
      logDiv.innerHTML = '';
      es = new EventSource(`/run?config=${qs}`);
      document.getElementById('stop').disabled = false;
      es.onmessage = (ev) => {
        try {
          const data = JSON.parse(ev.data);
          if (data.type === 'server_check') {
            log(`[check] ${data.url}: ${data.ok ? 'ok' : 'not responding'}`, data.ok ? 'ok':'err');
          } else if (data.type === 'round_start') {
            log(`===== RUNDA ${data.round} =====`);
          } else if (data.type === 'machine_response') {
            log(`[${data.rol} @ ${data.url} :: tokens=${data.tokens}] idei noi: ${data.new_ideas}`);
            log(data.text);
          } else if (data.type === 'round_end') {
            log(`-- sfârșit runda ${data.round}, pauză ${data.pause_sec}s --`);
          } else if (data.type === 'error') {
            log(`Eroare: ${data.url} - ${data.message}`, 'err');
          } else if (data.type === 'done') {
            log(data.message);
          }
        } catch (e) {
          log(`Parse error: ${e}`, 'err');
        }
      };
      es.onerror = () => {
        log('Conexiune întreruptă.', 'err');
        es.close();
      };
    };

    document.getElementById('stop').onclick = () => {
      if (es) es.close();
      document.getElementById('stop').disabled = true;
      log('Oprit de utilizator.');
    };
  </script>
</body>
</html>
    """
    return html


@app.get("/run")
def run_sse():
    cfg_raw = request.args.get("config")
    try:
        config = json.loads(cfg_raw) if cfg_raw else {}
    except Exception:
        config = {}
    return Response(stream_with_context(run_conversation(config)), mimetype="text/event-stream")


if __name__ == "__main__":
    # Rulează Flask pe localhost:5000
    app.run(host="127.0.0.1", port=5000, debug=False)