import os
import re
import time
import threading
import queue
import requests
from flask import Flask, request, jsonify
from flask_cors import CORS
from playwright.sync_api import sync_playwright

EMAIL = os.environ.get("HEDGEDOC_EMAIL")
PASSWORD = os.environ.get("HEDGEDOC_PASSWORD")
BASE_URL = os.environ.get("HEDGEDOC_BASE_URL", "https://hedgedoc.sa6anw.se")

# LLM Configuration
LLM_API_KEY = os.environ.get("LLM_API_KEY")
LLM_MODEL = os.environ.get("LLM_MODEL", "openai/gpt-4o-mini")
LLM_API_BASE = "https://openrouter.ai/api/v1/chat/completions"

AGENT_BLOCK_RE = re.compile(r"<!--Agent\n(.*?)\n-->", re.DOTALL)
NOTES_FILE = "active_notes.txt"

# Shared Queues and Flask
new_notes_queue = queue.Queue()
app = Flask(__name__)

CORS(app)

# ---------- LLM Integration ----------
def call_llm(context_text, chat_history):
    if not LLM_API_KEY:
        return "Error: LLM_API_KEY not set"

    system_prompt = (
        "You are a helpful assistant integrated into a HedgeDoc markdown document.\n\n"
        "DOCUMENT CONTEXT (everything outside the chat block):\n"
        f"{context_text}\n\n"
        "Instructions: Use the document context to answer questions. Be concise and maintain markdown formatting. "
        "Each chat block in the document is independent."
    )
    
    messages = [{"role": "system", "content": system_prompt}]
    
    for line in chat_history.splitlines():
        if line.startswith("User: "):
            messages.append({"role": "user", "content": line[6:].strip()})
        elif line.startswith("Agent: "):
            messages.append({"role": "assistant", "content": line[7:].strip()})

    try:
        response = requests.post(
            LLM_API_BASE,
            headers={
                "Authorization": f"Bearer {LLM_API_KEY}",
                "Content-Type": "application/json"
            },
            json={
                "model": LLM_MODEL,
                "messages": messages
            },
            timeout=60
        )
        response.raise_for_status()
        return response.json()['choices'][0]['message']['content']
    except Exception as e:
        return f"Error calling LLM: {e}"

# ---------- Data Persistence ----------
def load_persisted_notes():
    if not os.path.exists(NOTES_FILE):
        return []
    with open(NOTES_FILE, "r") as f:
        return list(set(line.strip() for line in f if line.strip()))

def persist_note(url):
    existing = load_persisted_notes()
    if url not in existing:
        with open(NOTES_FILE, "a") as f:
            f.write(f"{url}\n")

def remove_note_from_persistence(url):
    existing = load_persisted_notes()
    if url in existing:
        with open(NOTES_FILE, "w") as f:
            for line in existing:
                if line != url:
                    f.write(f"{line}\n")

# ---------- Web API ----------
@app.route('/')
def index():
    from flask import render_template
    return render_template('index.html')

@app.route('/add_note', methods=['POST'])
def add_note():
    data = request.get_json()
    if not data or 'note_id' not in data:
        return jsonify({"error": "Missing note_id"}), 400
    
    note_id = data['note_id']
    url = note_id if note_id.startswith("http") else f"{BASE_URL}/{note_id}"
    if "?edit" not in url:
        url += "?edit"
    
    persist_note(url)
    new_notes_queue.put(url)
    return jsonify({"status": "added", "url": url}), 200

def start_web_server():
    app.run(host='0.0.0.0', port=5000, debug=False, use_reloader=False)

# ---------- Hedgedoc Logic ----------
def login(page):
    try:
        print(f"[+] Logging in to {BASE_URL}...")
        page.goto(BASE_URL, wait_until="domcontentloaded", timeout=30000)
        sign_in_btn = page.wait_for_selector('button:has-text("Sign in")', timeout=10000)
        if sign_in_btn:
            sign_in_btn.click()
        page.wait_for_selector('input[name="email"]', timeout=5000)
        page.fill('input[name="email"]', EMAIL)
        page.fill('input[name="password"]', PASSWORD)
        page.click('form button[type="submit"]')
        page.wait_for_load_state("networkidle")
        print("[✓] Login successful")
    except Exception as e:
        print(f"[!] Login failed: {e}")

def is_logged_in(page):
    try:
        sign_in_visible = page.is_visible('button:has-text("Sign in")', timeout=2000)
        return not sign_in_visible
    except:
        return False

def ensure_logged_in(context):
    temp_page = context.new_page()
    try:
        temp_page.goto(BASE_URL, wait_until="domcontentloaded", timeout=10000)
        if not is_logged_in(temp_page):
            print("[!] Session expired or not logged in. Re-logging...")
            login(temp_page)
            return True
        return False
    except Exception as e:
        print(f"[!] Error checking login status: {e}")
        return False
    finally:
        temp_page.close()

def inject_event_bridge(page):
    try:
        page.evaluate("""
            () => {
                if (window.__agentBridgeInstalled) return;
                window.__agentBridgeInstalled = true;
                window.__agentEvents = [];
                const cm = document.querySelector('.CodeMirror').CodeMirror;
                cm.on("changes", (cm, changes) => {
                    window.__agentEvents.push({ ts: Date.now(), value: cm.getValue() });
                });
            }
        """)
    except Exception as e:
        print(f"[!] Failed to inject event bridge: {e}")

def ensure_agent_block(page):
    try:
        text = page.evaluate("() => document.querySelector('.CodeMirror').CodeMirror.getValue()")
        if "<!--Agent" not in text:
            print(f"[+] Creating initial Agent block on {page.url}")
            page.evaluate("""
                () => {
                    const cm = document.querySelector('.CodeMirror').CodeMirror;
                    cm.replaceRange("\\n\\n<!--Agent\\n-> \\n-->", { line: cm.lineCount(), ch: 0 });
                }
            """)
    except Exception as e:
        print(f"[!] Error in ensure_agent_block: {e}")

def find_prompts(text: str):
    prompts = []
    context_parts = AGENT_BLOCK_RE.split(text)
    full_context = "\n".join(part.strip() for part in context_parts if part.strip())

    for m in AGENT_BLOCK_RE.finditer(text):
        inner = m.group(1)
        lines = inner.splitlines()
        if not lines: continue

        for i in range(len(lines) - 1, -1, -1):
            if lines[i].strip().startswith("->"):
                if i + 1 < len(lines) and lines[-1].strip() == "":
                    prompt_part = lines[i][2:].strip()
                    extra_lines = lines[i+1:-1]
                    full_prompt = "\n".join([prompt_part] + extra_lines).strip()
                    
                    if full_prompt:
                        prompts.append({
                            "context": full_context,
                            "history": "\n".join(lines[:i]),
                            "prompt": full_prompt,
                            "start": m.start(),
                            "end": m.end(),
                            "original_lines": lines[:i]
                        })
                break
    return prompts

# ---------- Main Loop ----------
def main():
    if not EMAIL or not PASSWORD:
        print("[!] Error: HEDGEDOC_EMAIL and HEDGEDOC_PASSWORD must be set.")
        return

    threading.Thread(target=start_web_server, daemon=True).start()
    print("[✓] Web API listening on port 5000")

    with sync_playwright() as p:
        browser = p.chromium.launch(headless=True)
        context = browser.new_context()
        
        auth_page = context.new_page()
        login(auth_page)
        auth_page.close()
        
        pages = []
        def setup_note_page(url):
            try:
                print(f"[+] Opening note: {url}")
                page = context.new_page()
                page.goto(url, wait_until="networkidle", timeout=30000)
                page.keyboard.press("Control+Alt+E")
                page.wait_for_selector(".CodeMirror", timeout=15000)
                inject_event_bridge(page)
                ensure_agent_block(page)
                return page
            except Exception as e:
                print(f"[!] Failed to track {url}: {e}")
                return None

        for url in load_persisted_notes():
            pg = setup_note_page(url)
            if pg: pages.append(pg)

        print(f"[✓] Monitoring {len(pages)} notes. Ready.")
        last_login_check = time.time()

        while True:
            if time.time() - last_login_check > 600:
                if ensure_logged_in(context):
                    print("[✓] Re-authentication successful.")
                last_login_check = time.time()

            try:
                # 1. Handle new notes queue
                try:
                    while True:
                        url = new_notes_queue.get_nowait()
                        if any(p.url.split('?')[0] == url.split('?')[0] for p in pages if not p.is_closed()):
                            continue
                        pg = setup_note_page(url)
                        if pg: pages.append(pg)
                except queue.Empty:
                    pass

                # 2. Process active pages
                active_pages = []
                any_activity = False
                for page in pages:
                    if page.is_closed(): continue
                    
                    try:
                        content = page.evaluate("() => document.querySelector('.CodeMirror').CodeMirror.getValue()")
                    except Exception:
                        print(f"[!] Cannot access page {page.url}, checking session...")
                        if not is_logged_in(page):
                            ensure_logged_in(context)
                            page.reload(wait_until="networkidle")
                            continue
                        active_pages.append(page)
                        continue

                    if "<!--Agent" not in content:
                        print(f"[-] Releasing {page.url}")
                        remove_note_from_persistence(page.url)
                        page.close()
                        continue
                    
                    active_pages.append(page)
                    
                    num_events = page.evaluate("() => (window.__agentEvents || []).length")
                    if num_events and num_events > 0:
                        any_activity = True
                        event = page.evaluate("() => window.__agentEvents.shift()")
                        
                        prompts = find_prompts(event["value"])
                        for p in prompts:
                            # a. Show thinking state
                            thinking_text = f"<!--Agent\n" + "\n".join(p["original_lines"] + [f"User: {p['prompt']}", "Agent: 🧠 Tänker..."]) + "\n-->"
                            page.evaluate("""
                                (up) => {
                                    const cm = document.querySelector('.CodeMirror').CodeMirror;
                                    const doc = cm.getDoc();
                                    const from = doc.posFromIndex(up.start);
                                    const to = doc.posFromIndex(up.end);
                                    doc.replaceRange(up.new_text, from, to);
                                }
                            """, {"start": p["start"], "end": p["end"], "new_text": thinking_text})
                            
                            # b. Call LLM asychronously
                            response_queue = queue.Queue()
                            def llm_worker():
                                try:
                                    res = call_llm(p["context"], p["history"] + f"\nUser: {p['prompt']}")
                                    response_queue.put(("success", res))
                                except Exception as e:
                                    response_queue.put(("error", str(e)))
                            threading.Thread(target=llm_worker, daemon=True).start()
                            
                            # c. Animation loop
                            dots = ""
                            start_time = time.time()
                            last_dot_time = start_time
                            current_block_text = thinking_text
                            
                            while True:
                                try:
                                    status, ai_response = response_queue.get(timeout=0.5)
                                    break
                                except queue.Empty:
                                    now = time.time()
                                    if now - start_time > 60:
                                        status, ai_response = "error", "Timeout: Modellen svarade inte i tid."
                                        break
                                    if now - last_dot_time >= 5:
                                        dots += "."
                                        last_dot_time = now
                                        new_thinking = f"<!--Agent\n" + "\n".join(p["original_lines"] + [f"User: {p['prompt']}", f"Agent: 🧠 Tänker{dots}"]) + "\n-->"
                                        page.evaluate("""
                                            (up) => {
                                                const cm = document.querySelector('.CodeMirror').CodeMirror;
                                                const doc = cm.getDoc();
                                                const from = doc.posFromIndex(up.start);
                                                const fullText = doc.getValue();
                                                const endIdx = fullText.indexOf("-->", up.start);
                                                const to = (endIdx !== -1) ? doc.posFromIndex(endIdx + 3) : doc.posFromIndex(up.start + up.oldLen);
                                                doc.replaceRange(up.new_text, from, to);
                                            }
                                        """, {"start": p["start"], "oldLen": len(current_block_text), "new_text": new_thinking})
                                        current_block_text = new_thinking

                            # d. Show final response
                            if status == "error":
                                ai_response = f"❌ Fel: {ai_response}"
                                
                            final_text = f"<!--Agent\n" + "\n".join(p["original_lines"] + [f"User: {p['prompt']}", f"Agent: {ai_response}", "-> "]) + "\n-->"
                            page.evaluate("""
                                (up) => {
                                    const cm = document.querySelector('.CodeMirror').CodeMirror;
                                    const doc = cm.getDoc();
                                    const from = doc.posFromIndex(up.start);
                                    const fullText = doc.getValue();
                                    const endIdx = fullText.indexOf("-->", up.start);
                                    const to = (endIdx !== -1) ? doc.posFromIndex(endIdx + 3) : doc.posFromIndex(up.start + up.oldLen);
                                    doc.replaceRange(up.new_text, from, to);
                                    
                                    const lines = up.new_text.split('\\n');
                                    let foundLine = -1;
                                    for(let i=lines.length-1; i>=0; i--) {
                                        if(lines[i].includes("-> ")) {
                                            foundLine = from.line + i;
                                            break;
                                        }
                                    }
                                    if (foundLine !== -1) {
                                        cm.setCursor({line: foundLine, ch: 3});
                                    }
                                    cm.focus();
                                }
                            """, {"start": p["start"], "oldLen": len(current_block_text), "new_text": final_text})

                pages = active_pages
                if not any_activity:
                    time.sleep(0.5)

            except Exception as e:
                print(f"[!] Loop error: {e}")
                time.sleep(1)

if __name__ == "__main__":
    main()