import json, subprocess, threading, re, os

_processes = {}

def _update_progress(proc_id : int):
    try:
        proc = _processes[proc_id]["process"]
        while True:
            line_bytes = proc.stdout.readline()
            if not line_bytes and proc.poll() is not None:
                break
            
            if line_bytes:
                # Use replace to handle any weird characters gracefully
                line = line_bytes.decode('utf-8', errors='replace').strip()
                if line:
                    _processes[proc_id]["log"].append(line)
                    
                    # Match progress like [#abc 1.2MiB/10MiB(12%) CN:1 SD:2 DL:576KiB UL:1MiB]
                    prog_match = re.search(r"\[#\w+\s+([\w.]+)/([\w.]+)(?:\((\d+)%\))?.*?CN:(\d+).*?(?:SD:(\d+).*?)?DL:([\w.]+)(?:.*?UL:([\w.]+))?", line)
                    if prog_match:
                        _processes[proc_id]["dl_now"] = prog_match.group(1)
                        _processes[proc_id]["dl_total"] = prog_match.group(2)
                        _processes[proc_id]["progress"] = int(prog_match.group(3)) if prog_match.group(3) else 0
                        _processes[proc_id]["cn"] = prog_match.group(4)
                        _processes[proc_id]["sd"] = prog_match.group(5) if prog_match.group(5) else "0"
                        _processes[proc_id]["dl_speed"] = prog_match.group(6)
                        _processes[proc_id]["ul_speed"] = prog_match.group(7) if prog_match.group(7) else "0B"
                    
                    file_match = re.search(r"^FILE:\s+(.*)", line)
                    if file_match:
                        _processes[proc_id]["filename"] = os.path.basename(file_match.group(1))

        proc.wait()
        
        # Don't overwrite status if it was already marked canceled
        if _processes[proc_id].get("status") != "canceled":
            _processes[proc_id]["status"] = "success" if proc.returncode == 0 else "error"
        
        if proc.returncode != 0 and _processes[proc_id].get("status") != "canceled":
            _processes[proc_id]["log"].append(f"aria2c exited with code {proc.returncode}")
        print(f"-- Finished aria2 download ID {proc_id} with return code {proc.returncode}")
    except Exception as e:
        if _processes.get(proc_id, {}).get("status") != "canceled":
            _processes[proc_id]["status"] = "error"
            _processes[proc_id]["log"].append(f"Python exception in progress thread: {str(e)}")
        print(f"-- Exception in aria2 thread {proc_id}: {e}")

def main(args : dict):
    try:
        cmd = json.loads(args["txt"])
    except Exception as e:
        print(f"Error parsing args: {e}")
        return {"stdout": json.dumps({"error": "invalid json args"})}
    
    action = cmd.get("action")

    if action == "aria2.list":
        active_ids = [pid for pid, p in _processes.items() if p["user"] == args.get("user") and p["status"] == "running"]
        return {"stdout": json.dumps({"active_ids": active_ids})}

    if action == "aria2.cancel" and "id" in cmd:
        try:
            proc_id = int(cmd.get("id"))
            proc = _processes.get(proc_id)
            if proc and proc["user"] == args.get("user"):
                proc["status"] = "canceled"
                proc["process"].terminate()
                return {"stdout": json.dumps({"status": "canceled"})}
            return {"stdout": json.dumps({"error": "not found"})}
        except Exception as e:
            return {"stdout": json.dumps({"error": f"cancel check failed: {str(e)}"})}

    if action == "aria2.download" and "url" in cmd:
        if "w" not in args.get("perms", ""):
            return {"stdout": "error: no write permission"}
        
        url = cmd.get("url")
        target_dir = args.get("ap")
        filename = cmd.get("filename")

        # Disable console readout to avoid \r characters breaking readline()
        proc_cmd = [
            "aria2c", "--show-console-readout=false", "--summary-interval=1", url, "--dir", target_dir, 
            "-c", "-x16", "-s16", "-k1M", "--min-split-size=1M", "--file-allocation=none", 
            "--bt-max-peers=150", "--listen-port=50000-55000", "--dht-listen-port=50000-55000", "--enable-dht=true", "--enable-peer-exchange=true",
            "--bt-prioritize-piece=head=20M,tail=20M"
        ]
        if filename:
            proc_cmd += ["--out", filename]
        
        proc_id = len(_processes)
        try:
            _processes[proc_id] = {
                "process": subprocess.Popen(proc_cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT),
                "progress": 0,
                "status": "running",
                "log": [],
                "user": args.get("user"),
                "dl_now": "0",
                "dl_total": "Unknown",
                "dl_speed": "0",
                "filename": filename if filename else "Unknown"
            }
            threading.Thread(target=_update_progress, args=(proc_id, )).start()
            print(f"-- Started aria2 download ID {proc_id} for {url} to {target_dir}")
            return {"stdout": str(proc_id)}
        except Exception as e:
            print(f"Failed to start aria2c: {e}")
            return {"stdout": "error: failed to start aria2c"}
    
    if action == "aria2.status" and "id" in cmd:
        try:
            proc_id = int(cmd.get("id"))
            proc = _processes.get(proc_id)
            if not proc or proc["user"] != args.get("user"):
                return {"stdout": json.dumps({"error": "can't access download"})}
            
            resp = {
                "progress": proc["progress"],
                "status": proc["status"],
                "dl_now": proc.get("dl_now", "0"),
                "dl_total": proc.get("dl_total", "Unknown"),
                "dl_speed": proc.get("dl_speed", "0"),
                "ul_speed": proc.get("ul_speed", "0B"),
                "cn": proc.get("cn", "0"),
                "sd": proc.get("sd", "0"),
                "filename": proc.get("filename", "Unknown"),
                "log": proc["log"][-10:] # Return last 10 lines
            }

            if proc["status"] != "running":
                _processes.pop(proc_id)
            
            return {"stdout": json.dumps(resp)}
        except Exception as e:
            return {"stdout": json.dumps({"error": f"status check failed: {str(e)}"})}

    return {"stdout": json.dumps({"error": "unknown action"})}