app.py

import os
import sys
import time
import subprocess
import urllib.request
import socket
import gradio as gr
import http.client
import urllib.parse
import xml.etree.ElementTree as ET
import shutil
import re
import glob

# =============================================================================
# CONFIGURATION
# =============================================================================
REPO_URL = "https://github.com/Automattic/atd-server-next.git"
SERVER_DIR = "atd-server-next"
MODELS_DIR = os.path.join(SERVER_DIR, "models")
MODEL_BASE_URL = "https://openatd.svn.wordpress.org/atd-server/models/"
HOST = "127.0.0.1"
PORT = 1049

MODEL_FILES = [
    "cnetwork.bin", "cnetwork2.bin", "dictionary.txt", "edits.bin",
    "endings.bin", "hnetwork.bin", "hnetwork2.bin", "hnetwork4.bin",
    "lexicon.bin", "model.bin", "model.zip", "network3f.bin",
    "network3p.bin", "not_misspelled.txt", "stringpool.bin", "trigrams.bin"
]

# =============================================================================
# PHASE 1: SETUP & PATCHING
# =============================================================================
def recursive_patch_code():
    """
    Aggressively finds AND removes incompatible imports in ALL .sl files.
    The syntax 'from: jarfile.jar' breaks on Java 9+.
    We replace it with nothing, relying on the classpath passed to the java command.
    """
    print("--- [PHASE 1.5] RECURSIVE CODE PATCHING (JAVA 17 COMPATIBILITY) ---")
    
    # Find all .sl files in the server directory
    sl_files = glob.glob(f"{SERVER_DIR}/**/*.sl", recursive=True)
    
    patched_count = 0
    # Regex explanation:
    # from:     Matches the keyword
    # \s* Matches whitespace
    # [^;\n]+   Matches anything that ISN'T a semicolon or newline (the jar path)
    # (?=;)     Lookahead to ensure we stop at the semicolon
    regex = re.compile(r'from:\s*[^;\n]+(?=;)', re.IGNORECASE)

    for file_path in sl_files:
        try:
            with open(file_path, "r", encoding="utf-8", errors="ignore") as f:
                content = f.read()
            
            # Apply the regex
            new_content = regex.sub('', content)
            
            if content != new_content:
                print(f"Patching {file_path}...")
                with open(file_path, "w", encoding="utf-8") as f:
                    f.write(new_content)
                patched_count += 1
        except Exception as e:
            print(f"Warning: Could not patch {file_path}: {e}")
            
    print(f"-> Successfully patched {patched_count} files.")

def setup_server():
    print("--- [PHASE 0] CHECKING REPOSITORY ---")
    if not os.path.exists(SERVER_DIR):
        print(f"Repository not found. Cloning from {REPO_URL}...")
        subprocess.run(["git", "clone", "--depth", "1", REPO_URL, SERVER_DIR], check=True)

    print("\n--- [PHASE 1] CHECKING MODELS ---")
    if not os.path.exists(MODELS_DIR):
        os.makedirs(MODELS_DIR, exist_ok=True)

    for filename in MODEL_FILES:
        filepath = os.path.join(MODELS_DIR, filename)
        if not os.path.exists(filepath):
            url = MODEL_BASE_URL + filename
            print(f"Downloading {filename}...")
            try:
                urllib.request.urlretrieve(url, filepath)
            except Exception as e:
                print(f"  -> FAILED: {e}")
                sys.exit(1)

    # Patch code BEFORE compiling
    recursive_patch_code()

    print("\n--- [PHASE 2] COMPILING RULES ---")
    # We compile if rules.bin is missing. 
    # Note: The compiler output is actually rules.bin (or parts of it loaded into memory)
    # The previous error showed FileNotFound for models/rules.bin.
    
    # Classpath separator: ':' for Linux (Hugging Face), ';' for Windows
    cp_sep = ";" if os.name == 'nt' else ":"
    
    try:
        # We explicitly add the JARs to the compiler classpath too
        cp = f"lib/sleep.jar{cp_sep}lib/moconti.jar{cp_sep}lib/spellutils.jar"
        
        cmd = [
            "java", 
            "-Datd.lowmem=true", 
            "-Xmx1024M", 
            "-classpath", cp,
            "sleep.console.TextConsole", # Run the console directly to avoid jar limits
            "utils/rules/rules.sl"
        ]
        
        print("Executing Rule Compiler...")
        subprocess.run(cmd, cwd=SERVER_DIR, check=True)
        print("Rules compiled successfully.")
        
    except subprocess.CalledProcessError as e:
        print(f"\nFATAL ERROR: Rule compilation failed with code {e.returncode}.")
        print("This usually means the code patch failed or Java is incompatible.")
        sys.exit(1)

# =============================================================================
# PHASE 2: SERVER RUNNER
# =============================================================================
def start_backend():
    print("\n--- [PHASE 3] STARTING SERVER ---")
    cp_sep = ";" if os.name == 'nt' else ":"
    
    # Robust Classpath
    classpath = f"lib/sleep.jar{cp_sep}lib/moconti.jar{cp_sep}lib/spellutils.jar"
    sleep_cp = f"lib{cp_sep}service/code"

    cmd = [
        "java",
        "-Dfile.encoding=UTF-8",
        "-XX:+AggressiveHeap",
        "-XX:+UseParallelGC",
        "-Datd.lowmem=true",
        "-Dbind.interface=127.0.0.1",
        f"-Dserver.port={PORT}",
        f"-Dsleep.classpath={sleep_cp}",
        "-Dsleep.debug=24",
        "-classpath", classpath,
        "httpd.Moconti",
        "atdconfig.sl"
    ]
    
    print("Launching Java subprocess...")
    return subprocess.Popen(cmd, cwd=SERVER_DIR)

def wait_for_port(timeout=60):
    print(f"Waiting for port {PORT}...")
    start = time.time()
    while time.time() - start < timeout:
        try:
            with socket.create_connection((HOST, PORT), timeout=1):
                print("Server is Online!")
                return True
        except (ConnectionRefusedError, OSError):
            time.sleep(1)
    return False

# =============================================================================
# PHASE 3: CLIENT
# =============================================================================
class AtDClient:
    def check_document(self, text):
        try:
            conn = http.client.HTTPConnection(HOST, PORT, timeout=5)
            params = urllib.parse.urlencode({'key': 'gradio', 'data': text})
            headers = {"Content-Type": "application/x-www-form-urlencoded"}
            conn.request("POST", "/checkDocument", params, headers)
            resp = conn.getresponse()
            
            # If server is still warming up, it might return 503 or connection drop
            if resp.status != 200: return []
            
            # Read response
            xml_text = resp.read().decode('utf-8', errors='ignore')
            
            # Basic XML validation
            if not xml_text.strip().startswith("<"): return []

            root = ET.fromstring(xml_text)
            errors = []
            for e in root.findall('error'):
                err = {
                    'string': e.find('string').text,
                    'description': e.find('description').text,
                    'type': e.find('type').text,
                    'precontext': e.find('precontext').text or "",
                    'suggestions': []
                }
                sug = e.find('suggestions')
                if sug is not None:
                    err['suggestions'] = [o.text for o in sug.findall('option') if o.text]
                errors.append(err)
            return errors
        except Exception as e:
            print(f"Client Error: {e}")
            return []

client = AtDClient()

def analyze_text(text):
    if not text.strip(): return []
    errors = client.check_document(text)
    output = []
    last_pos = 0
    
    for err in errors:
        word = err['string']
        search_start = last_pos
        if err['precontext']:
            context_idx = text.find(err['precontext'], last_pos)
            if context_idx != -1:
                search_start = context_idx + len(err['precontext'])

        idx = text.find(word, search_start)
        if idx != -1:
            if idx > last_pos:
                output.append((text[last_pos:idx], None))
            
            label = f"{err['type']}: {err['description']}"
            if err['suggestions']:
                label += f" -> {', '.join(err['suggestions'][:3])}"
            
            output.append((text[idx:idx+len(word)], label))
            last_pos = idx + len(word)
            
    if last_pos < len(text):
        output.append((text[last_pos:], None))
    return output

if __name__ == "__main__":
    setup_server()
    server_proc = start_backend()
    
    # Monitoring loop
    # We check the process status before starting the UI
    time.sleep(3)
    if server_proc.poll() is not None:
        print("FATAL: Java server exited immediately. Check console for ClassCastException.")
        sys.exit(1)

    if wait_for_port(timeout=120): # Give it 2 minutes to load models
        with gr.Blocks(title="AtD Self-Hosted") as demo:
            gr.Markdown("# 🛡️ After The Deadline")
            gr.Markdown("Running on Java 17 (Patched)")
            
            with gr.Row():
                inp = gr.Textbox(label="Input", placeholder="Type text here...", lines=6)
                out = gr.HighlightedText(label="Analysis", combine_adjacent=True)
            
            btn = gr.Button("Check Text", variant="primary")
            btn.click(analyze_text, inputs=inp, outputs=out)
            
        demo.launch(server_name="0.0.0.0", server_port=7860)
    else:
        print("FATAL: Server did not start (Timeout).")
        server_proc.kill()