import http.server, socketserver, json, base64, os, re, sys, argparse
from datetime import datetime

def parse_args():
    p = argparse.ArgumentParser(description="NovelAI autosave listener")
    p.add_argument('--port', type=int, default=8891)
    p.add_argument('--dir', default=r"A:\sd-webui-forge-neo\outputs")
    return p.parse_args()

args = parse_args()
PORT = args.port
SAVE_DIR = args.dir
os.makedirs(SAVE_DIR, exist_ok=True)

INVALID_CHARS = re.compile(r'[<>:"/\\|?*\x00-\x1f]')
COUNTER_RE = re.compile(r'^(\d+)-')
_counters: dict[str, int] = {}
port = 18891

def log(msg):
    print(f"[{datetime.now().strftime('%H:%M:%S %d-%m-%Y')}] {msg}")

def sanitize(text, max_len=100):
    text = INVALID_CHARS.sub('_', text).strip(' .')
    return text[:max_len] or 'untitled'

def next_counter(subfolder):
    if subfolder not in _counters:
        folder = os.path.join(SAVE_DIR, subfolder)
        os.makedirs(folder, exist_ok=True)
        nums = [int(m.group(1)) for f in os.listdir(folder) if (m := COUNTER_RE.match(f))]
        _counters[subfolder] = max(nums, default=0)
    _counters[subfolder] += 1
    return _counters[subfolder]

def fmt_num(n):
    if isinstance(n, float) and n.is_integer():
        return str(int(n))
    return str(n)

def format_character_prompts(comment):
    pos_root = comment.get('v4_prompt') or {}
    neg_root = comment.get('v4_negative_prompt') or {}
    pos_chars = ((pos_root.get('caption') or {}).get('char_captions')) or []
    neg_chars = ((neg_root.get('caption') or {}).get('char_captions')) or []
    use_coords = pos_root.get('use_coords', False)

    n = max(len(pos_chars), len(neg_chars))
    if n == 0:
        return []

    lines = []
    for i in range(n):
        pos_entry = pos_chars[i] if i < len(pos_chars) else {}
        neg_entry = neg_chars[i] if i < len(neg_chars) else {}
        pos_text = pos_entry.get('char_caption', '')
        neg_text = neg_entry.get('char_caption', '')
        if not pos_text and not neg_text:
            continue

        pos_str = ''
        if use_coords:
            centers = pos_entry.get('centers') or neg_entry.get('centers') or []
            if centers:
                c = centers[0]
                pos_str = f" (pos: {fmt_num(c.get('x'))}, {fmt_num(c.get('y'))})"

        label = f"Character {i + 1}{pos_str}:"
        lines.append(f"{label} {pos_text}" if pos_text else f"{label} (no caption)")
        if neg_text:
            lines.append(f"Character {i + 1} negative: {neg_text}")
    return lines

def build_plaintext_metadata(comment):
    prompt = comment.get('prompt', '')
    negative = comment.get('uc', '')
    schedule = str(comment.get('noise_schedule', '')).capitalize()

    params = [
        f"Steps: {comment.get('steps', '')}",
        f"Sampler: {comment.get('sampler', '')}",
        f"Schedule type: {schedule}",
        f"CFG scale: {fmt_num(comment.get('scale', ''))}",
        f"Seed: {comment.get('seed', '')}",
        f"Size: {comment.get('width', '')}x{comment.get('height', '')}",
        f"Model: {comment.get('model_name', '')}",
    ]
    if comment.get('model_hash'):
        params.append(f"Model hash: {comment['model_hash']}")

    lines = [prompt, f"Negative prompt: {negative}"]
    lines.extend(format_character_prompts(comment))
    lines.append(', '.join(params))
    return '\n'.join(lines)

class Handler(http.server.BaseHTTPRequestHandler):
    def do_GET(self):
        self.send_response(200)
        self.send_header('Access-Control-Allow-Origin', '*')
        self.send_header('Content-Type', 'text/plain')
        self.end_headers()
        self.wfile.write(b'ok')

    def do_POST(self):
        length = int(self.headers.get('Content-Length', 0))
        payload = json.loads(self.rfile.read(length))
        img_bytes = base64.b64decode(payload['data'])
        meta = payload.get('metadata', {})
        ext = payload.get('ext', 'png')
        if ext not in ('png', 'webp'):
            ext = 'png'

        subfolder = payload.get('subfolder', 'txt2img-images')
        if subfolder not in ('txt2img-images', 'img2img-images'):
            subfolder = 'txt2img-images'  # guard against unexpected values
        folder = os.path.join(SAVE_DIR, subfolder)
        os.makedirs(folder, exist_ok=True)

        try:
            comment = json.loads(meta.get('Comment', '{}'))
        except json.JSONDecodeError:
            comment = {}
        seed = comment.get('seed', 'noseed')
        prompt = meta.get('Description') or comment.get('prompt', '')

        n = next_counter(subfolder)
        base = f"{n}-{seed}-{sanitize(prompt)}"

        with open(os.path.join(folder, base + '.' + ext), 'wb') as f:
            f.write(img_bytes)
        with open(os.path.join(folder, base + '.txt'), 'w', encoding='utf-8') as f:
            f.write(build_plaintext_metadata(comment))

        log(f"saved {subfolder}/{base}.{ext}")
        self.send_response(200)
        self.send_header('Access-Control-Allow-Origin', '*')
        self.end_headers()

    def do_OPTIONS(self):
        self.send_response(200)
        self.send_header('Access-Control-Allow-Origin', '*')
        self.send_header('Access-Control-Allow-Methods', 'POST, OPTIONS')
        self.send_header('Access-Control-Allow-Headers', 'Content-Type')
        self.end_headers()

    def log_message(self, *a): pass

if __name__ == '__main__':
    try:
        with socketserver.TCPServer(("127.0.0.1", PORT), Handler) as httpd:
            log(f"listening on http://127.0.0.1:{PORT}, saving to {SAVE_DIR}")
            httpd.serve_forever()
    except OSError as e:
        log(f"Couldn't bind to port {PORT}: {e}")
        log(f"Something's already using it. Try: python {sys.argv[0]} --port {PORT + 1}")