from __future__ import annotations

import base64
import concurrent.futures
import hashlib
import json
import mimetypes
import os
import re
import shutil
import sys
import threading
import time
import urllib.error
import urllib.parse
import urllib.request
from email.message import Message
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path

sys.dont_write_bytecode = True


HOST = "127.0.0.1"
PORT = 17661
MIN_RANGE_SIZE = 8 * 1024 * 1024
CHUNK_SIZE = 4 * 1024 * 1024
MAX_CONNECTIONS = 8
MAX_RETRIES = 3
REQUEST_TIMEOUT = 30
LOG_PATH = Path(__file__).with_name("download_bridge.log")
JOBS_PATH = Path(__file__).with_name("download_jobs.json")
JOB_HISTORY_LIMIT = 100
RUNNING_PROGRESS_KEYS = {"status", "downloaded", "total"}


download_pool = concurrent.futures.ThreadPoolExecutor(max_workers=4)
log_lock = threading.Lock()
jobs: dict[str, dict] = {}
jobs_lock = threading.Lock()
controls: dict[str, dict[str, bool]] = {}
controls_lock = threading.Lock()


class DownloadPaused(Exception):
    pass


class DownloadCancelled(Exception):
    pass


def default_download_dir() -> Path:
    configured = os.environ.get("ONE_DOWNLOAD_DIR")
    if configured:
        return Path(configured).expanduser()

    desktop = Path(os.environ.get("USERPROFILE", str(Path.home()))) / "Desktop"
    return desktop if desktop.exists() else Path.home()


DOWNLOAD_DIR = default_download_dir()


def log(message: str) -> None:
    line = f"{time.strftime('%Y-%m-%d %H:%M:%S')} {message}\n"
    with log_lock:
        with LOG_PATH.open("a", encoding="utf-8") as handle:
            handle.write(line)


def load_jobs() -> None:
    if not JOBS_PATH.exists():
        return

    try:
        data = json.loads(JOBS_PATH.read_text(encoding="utf-8"))
        loaded = data.get("jobs", [])
        if not isinstance(loaded, list):
            return
        changed = False
        with jobs_lock:
            for job in loaded[-JOB_HISTORY_LIMIT:]:
                if isinstance(job, dict) and job.get("id") and job.get("status") != "done":
                    jobs[str(job["id"])] = job
                else:
                    changed = True
        if changed:
            save_jobs()
    except Exception as error:
        log(f"load jobs failed error={error}")


def save_jobs() -> None:
    try:
        with jobs_lock:
            items = sorted((to_persisted_job(job) for job in jobs.values()), key=lambda job: job.get("updatedAt", 0))[-JOB_HISTORY_LIMIT:]
        temp_path = JOBS_PATH.with_suffix(".json.tmp")
        temp_path.write_text(json.dumps({"jobs": items}, ensure_ascii=False, indent=2), encoding="utf-8")
        temp_path.replace(JOBS_PATH)
    except Exception as error:
        log(f"save jobs failed error={error}")


def to_persisted_job(job: dict) -> dict:
    return {key: value for key, value in job.items() if not key.startswith("_")}


def should_save_job(previous_status: str, changes: dict, pruned: bool) -> bool:
    if pruned:
        return True
    if not changes:
        return False
    next_status = str(changes.get("status") or previous_status or "")
    if next_status == "running" and previous_status == "running" and set(changes) <= RUNNING_PROGRESS_KEYS:
        return False
    return True


def make_job_id(payload: dict) -> str:
    seed = f"{time.time_ns()}|{payload.get('url', '')}|{threading.get_ident()}"
    return hashlib.sha1(seed.encode("utf-8", "ignore")).hexdigest()[:16]


def update_job(job_id: str, **changes) -> None:
    if not job_id:
        return

    should_save = False
    with jobs_lock:
        job = jobs.setdefault(job_id, {"id": job_id, "createdAt": time.time()})
        previous_status = str(job.get("status") or "")
        job.update(changes)
        job["updatedAt"] = time.time()
        pruned = False
        if len(jobs) > JOB_HISTORY_LIMIT:
            old_ids = sorted(jobs, key=lambda key: jobs[key].get("updatedAt", 0))[: len(jobs) - JOB_HISTORY_LIMIT]
            for old_id in old_ids:
                jobs.pop(old_id, None)
                pruned = True
        should_save = should_save_job(previous_status, changes, pruned)
    if should_save:
        save_jobs()


def list_jobs() -> list[dict]:
    with jobs_lock:
        visible = []
        for job in jobs.values():
            if job.get("status") == "done":
                continue
            item = {key: value for key, value in job.items() if not key.startswith("_")}
            visible.append(item)
        return sorted(visible, key=lambda job: job.get("updatedAt", 0), reverse=True)


def get_job(job_id: str) -> dict | None:
    with jobs_lock:
        job = jobs.get(job_id)
        return dict(job) if job else None


def remove_job(job_id: str) -> None:
    with jobs_lock:
        jobs.pop(job_id, None)
    clear_control(job_id)
    save_jobs()


def ensure_control(job_id: str) -> dict[str, bool]:
    with controls_lock:
        return controls.setdefault(job_id, {"pause": False, "cancel": False})


def set_control(job_id: str, key: str, value: bool) -> None:
    with controls_lock:
        state = controls.setdefault(job_id, {"pause": False, "cancel": False})
        state[key] = value


def clear_control(job_id: str) -> None:
    with controls_lock:
        controls.pop(job_id, None)


def check_control(job_id: str) -> None:
    if not job_id:
        return

    with controls_lock:
        state = controls.get(job_id) or {}
        cancel = state.get("cancel", False)
        pause = state.get("pause", False)

    if cancel:
        raise DownloadCancelled()
    if pause:
        raise DownloadPaused()


def sanitize_filename(name: str) -> str:
    cleaned = re.sub(r'[\\/:*?"<>|]', "_", name or "")
    cleaned = re.sub(r"[\r\n]+", " ", cleaned)
    cleaned = re.sub(r"\s+", " ", cleaned).strip(" .")
    return cleaned or f"download_{time.strftime('%Y%m%d_%H%M%S')}.bin"


def unique_path(directory: Path, filename: str) -> Path:
    directory.mkdir(parents=True, exist_ok=True)
    candidate = directory / sanitize_filename(filename)
    if not candidate.exists():
        return candidate

    stem = candidate.stem
    suffix = candidate.suffix
    for index in range(2, 10000):
        candidate = directory / f"{stem}_{index}{suffix}"
        if not candidate.exists():
            return candidate
    return directory / f"{stem}_{int(time.time())}{suffix}"


def resolve_final_path(payload: dict, filename: str) -> Path:
    final_path_value = str(payload.get("finalPath") or "").strip()
    if final_path_value:
        final_path = Path(os.path.expandvars(final_path_value)).expanduser()
        if final_path.is_absolute():
            final_path.parent.mkdir(parents=True, exist_ok=True)
            return final_path

    save_dir_value = str(payload.get("saveDir") or "").strip()
    save_dir = Path(os.path.expandvars(save_dir_value)).expanduser() if save_dir_value else DOWNLOAD_DIR
    if not save_dir.is_absolute():
        save_dir = DOWNLOAD_DIR
    return unique_path(save_dir, filename)


def filename_from_content_disposition(value: str) -> str:
    if not value:
        return ""

    message = Message()
    message["content-disposition"] = value
    filename = message.get_filename()
    if filename:
        return filename

    match = re.search(r"filename\*=UTF-8''([^;]+)", value, re.I)
    if match:
        return urllib.parse.unquote(match.group(1))

    match = re.search(r'filename="?([^";]+)"?', value, re.I)
    return match.group(1).strip() if match else ""


def filename_from_url(url: str) -> str:
    parsed = urllib.parse.urlparse(url)
    name = urllib.parse.unquote(Path(parsed.path).name)
    return name or ""


def pick_filename(payload: dict, headers: dict) -> str:
    for candidate in (
        payload.get("filename", ""),
        filename_from_content_disposition(headers.get("content-disposition", "")),
        filename_from_url(payload.get("url", "")),
    ):
        if str(candidate).strip():
            return sanitize_filename(str(candidate))

    content_type = headers.get("content-type", "").split(";", 1)[0].strip()
    extension = mimetypes.guess_extension(content_type) or ".bin"
    return f"download_{time.strftime('%Y%m%d_%H%M%S')}{extension}"


def normalize_headers(payload: dict) -> dict:
    headers = dict(payload.get("headers") or {})
    headers.pop("Host", None)
    headers.pop("Content-Length", None)
    headers.pop("Connection", None)
    headers.pop("Accept-Encoding", None)

    if payload.get("referrer") and "Referer" not in headers:
        headers["Referer"] = str(payload["referrer"])
    if "User-Agent" not in headers:
        headers["User-Agent"] = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 Chrome/145 Safari/537.36"
    if "Accept" not in headers:
        headers["Accept"] = "*/*"

    return headers


def decode_post_body(payload: dict) -> bytes | None:
    encoded = payload.get("postDataBase64") or ""
    if not encoded:
        return None
    return base64.b64decode(encoded)


def request_once(url: str, headers: dict, method: str = "GET", data: bytes | None = None):
    req = urllib.request.Request(url, data=data, headers=headers, method=method)
    return urllib.request.urlopen(req, timeout=REQUEST_TIMEOUT)


def probe(payload: dict, headers: dict, body: bytes | None) -> dict:
    url = payload["url"]
    method = str(payload.get("method") or "GET").upper()
    info = {"size": 0, "range": False, "headers": {}, "status": 0}

    methods = ["HEAD"] if method == "GET" else []
    methods.append(method)

    for probe_method in methods:
        probe_headers = dict(headers)
        probe_data = body if probe_method not in ("GET", "HEAD") else None
        if probe_method == "GET":
            probe_headers["Range"] = "bytes=0-0"
        try:
            with request_once(url, probe_headers, probe_method, probe_data) as response:
                response_headers = {key.lower(): value for key, value in response.headers.items()}
                info["status"] = response.status
                info["headers"] = response_headers
                size = parse_size(response_headers)
                if size:
                    info["size"] = size
                if response.status == 206 or response_headers.get("accept-ranges", "").lower() == "bytes":
                    info["range"] = True
                if probe_method == "GET":
                    response.read(1)
                return info
        except urllib.error.HTTPError as error:
            response_headers = {key.lower(): value for key, value in error.headers.items()}
            info["status"] = error.code
            info["headers"] = response_headers
            size = parse_size(response_headers)
            if size:
                info["size"] = size
            if error.code == 206:
                info["range"] = True
                return info
            if error.code in (405, 501):
                continue
        except Exception:
            continue

    return info


def parse_size(headers: dict) -> int:
    content_range = headers.get("content-range", "")
    match = re.search(r"/(\d+)$", content_range)
    if match:
        return int(match.group(1))

    content_length = headers.get("content-length", "")
    return int(content_length) if content_length.isdigit() else 0


def stream_to_file(response, path: Path, job_id: str = "") -> None:
    with path.open("wb") as handle:
        while True:
            check_control(job_id)
            chunk = response.read(1024 * 1024)
            if not chunk:
                break
            handle.write(chunk)


def verify_file(path: Path, expected_size: int = 0) -> bool:
    if not path.exists() or path.stat().st_size <= 0:
        return False
    if expected_size and path.stat().st_size != expected_size:
        return False
    return True


def download_single(payload: dict, headers: dict, body: bytes | None, path: Path, job_id: str = "") -> None:
    method = str(payload.get("method") or "GET").upper()
    data = body if method not in ("GET", "HEAD") else None
    with request_once(payload["url"], headers, method, data) as response:
        stream_to_file(response, path, job_id)


def download_range_part(url: str, headers: dict, start: int, end: int, path: Path, job_id: str = "") -> None:
    expected_size = end - start + 1
    if verify_file(path, expected_size):
        return

    path.unlink(missing_ok=True)
    temp_path = path.with_suffix(path.suffix + ".tmp")
    part_headers = dict(headers)
    part_headers["Range"] = f"bytes={start}-{end}"
    last_error: Exception | None = None

    for attempt in range(1, MAX_RETRIES + 1):
        try:
            check_control(job_id)
            temp_path.unlink(missing_ok=True)
            with request_once(url, part_headers, "GET", None) as response:
                if response.status != 206:
                    raise RuntimeError(f"Range request failed: {response.status}")
                stream_to_file(response, temp_path, job_id)

            if verify_file(temp_path, expected_size):
                temp_path.replace(path)
                return
            raise RuntimeError("Range part verification failed")
        except (DownloadPaused, DownloadCancelled):
            temp_path.unlink(missing_ok=True)
            raise
        except Exception as error:
            last_error = error
            log(f"part failed attempt={attempt} range={start}-{end} error={error}")
            temp_path.unlink(missing_ok=True)
            time.sleep(0.5 * attempt)

    raise RuntimeError(str(last_error) if last_error else "Range part failed")


def download_accelerated(payload: dict, headers: dict, size: int, path: Path, job_id: str = "") -> None:
    parts = []
    start = 0
    index = 0
    while start < size:
        end = min(size - 1, start + CHUNK_SIZE - 1)
        part_path = path.with_suffix(path.suffix + f".part{index:03d}")
        parts.append((start, end, part_path))
        start = end + 1
        index += 1

    with concurrent.futures.ThreadPoolExecutor(max_workers=min(MAX_CONNECTIONS, len(parts))) as pool:
        futures = [
            pool.submit(download_range_part, payload["url"], headers, start, end, part_path, job_id)
            for start, end, part_path in parts
        ]
        for future in concurrent.futures.as_completed(futures):
            future.result()
            if job_id:
                downloaded = sum(part_path.stat().st_size for _, _, part_path in parts if part_path.exists())
                update_job(job_id, status="running", downloaded=downloaded, total=size)

    with path.open("wb") as output:
        for _, _, part_path in parts:
            check_control(job_id)
            with part_path.open("rb") as part:
                shutil.copyfileobj(part, output, 1024 * 1024)
            part_path.unlink(missing_ok=True)


def cleanup_download_parts(temp_path: Path) -> None:
    temp_path.unlink(missing_ok=True)
    for part_path in temp_path.parent.glob(temp_path.name + ".part*"):
        part_path.unlink(missing_ok=True)
    for part_path in temp_path.parent.glob(temp_path.name + ".part*.tmp"):
        part_path.unlink(missing_ok=True)


def download_worker(payload: dict) -> None:
    job_id = str(payload.get("jobId") or make_job_id(payload))
    payload["jobId"] = job_id
    ensure_control(job_id)
    url = str(payload.get("url") or "")
    if not re.match(r"^https?://", url, re.I):
        log(f"reject invalid url: {url}")
        update_job(job_id, status="failed", url=url, error="invalid url")
        return

    payload["url"] = url
    update_job(job_id, status="running", url=url, source=payload.get("source", ""), _payload=payload)
    headers = normalize_headers(payload)
    body = decode_post_body(payload)
    info = probe(payload, headers, body)
    filename = pick_filename(payload, info.get("headers") or {})
    final_path = resolve_final_path(payload, filename)
    payload["finalPath"] = str(final_path)
    payload["saveDir"] = str(final_path.parent)
    temp_path = final_path.with_name(final_path.name + ".onepart")
    update_job(
        job_id,
        status="running",
        file=str(final_path),
        filename=final_path.name,
        total=int(info.get("size") or 0),
        range=bool(info.get("range")),
        downloaded=0,
        _payload=payload,
    )

    for attempt in range(1, MAX_RETRIES + 1):
        try:
            temp_path.unlink(missing_ok=True)
            if (
                str(payload.get("method") or "GET").upper() == "GET"
                and info.get("range")
                and info.get("size", 0) >= MIN_RANGE_SIZE
            ):
                log(f"accelerated attempt={attempt} url={url} file={final_path}")
                download_accelerated(payload, headers, int(info["size"]), temp_path, job_id)
            else:
                log(f"single attempt={attempt} url={url} file={final_path}")
                download_single(payload, headers, body, temp_path, job_id)

            expected = int(info["size"]) if info.get("range") else 0
            if verify_file(temp_path, expected):
                temp_path.replace(final_path)
                log(f"done file={final_path}")
                remove_job(job_id)
                return
            raise RuntimeError("file verification failed")
        except DownloadPaused:
            log(f"paused url={url}")
            update_job(job_id, status="paused", error="", attempt=attempt)
            temp_path.unlink(missing_ok=True)
            return
        except DownloadCancelled:
            log(f"cancelled url={url}")
            cleanup_download_parts(temp_path)
            update_job(job_id, status="cancelled", error="", attempt=attempt)
            clear_control(job_id)
            return
        except Exception as error:
            log(f"failed attempt={attempt} url={url} error={error}")
            update_job(job_id, status="retrying", error=str(error), attempt=attempt)
            temp_path.unlink(missing_ok=True)
            time.sleep(0.5 * attempt)

    log(f"give up url={url}")
    update_job(job_id, status="failed", error="max retries exceeded")
    clear_control(job_id)


def submit_download(payload: dict, status: str = "queued") -> str:
    job_id = str(payload.get("jobId") or make_job_id(payload))
    payload["jobId"] = job_id
    set_control(job_id, "pause", False)
    set_control(job_id, "cancel", False)
    update_job(job_id, status=status, url=payload.get("url", ""), source=payload.get("source", ""), _payload=payload)
    download_pool.submit(download_worker, dict(payload))
    return job_id


def control_job(action: str, job_id: str) -> dict:
    job = get_job(job_id)
    if not job:
        return {"ok": False, "error": "작업을 찾지 못했습니다."}

    if action == "pause":
        set_control(job_id, "pause", True)
        update_job(job_id, status="pausing")
        return {"ok": True}

    if action == "cancel":
        set_control(job_id, "cancel", True)
        update_job(job_id, status="cancelling")
        return {"ok": True}

    if action in ("resume", "retry"):
        payload = job.get("_payload")
        if not isinstance(payload, dict):
            return {"ok": False, "error": "재시도할 작업 정보가 없습니다."}
        payload = dict(payload)
        payload["jobId"] = job_id
        submit_download(payload, "queued")
        return {"ok": True}

    if action == "clear":
        remove_job(job_id)
        return {"ok": True}

    return {"ok": False, "error": "지원하지 않는 작업입니다."}


class Handler(BaseHTTPRequestHandler):
    server_version = "ONEDownloadBridge/1.0"

    def do_OPTIONS(self) -> None:
        self.send_response(204)
        self.send_cors()
        self.end_headers()

    def do_GET(self) -> None:
        if self.path == "/health":
            self.send_json({"ok": True, "name": "ONE Download Bridge"})
            return
        path = urllib.parse.urlparse(self.path).path
        if path == "/config":
            self.send_json({"ok": True, "defaultDownloadDir": str(DOWNLOAD_DIR)})
            return
        if path == "/status":
            self.send_json({"ok": True, "jobs": list_jobs()})
            return
        self.send_error(404)

    def do_POST(self) -> None:
        path = urllib.parse.urlparse(self.path).path
        if path == "/control":
            self.handle_control()
            return
        if path != "/download":
            self.send_error(404)
            return

        try:
            length = int(self.headers.get("Content-Length", "0"))
            payload = json.loads(self.rfile.read(length).decode("utf-8"))
            job_id = submit_download(payload)
            self.send_json({"ok": True, "jobId": job_id})
        except Exception as error:
            log(f"post error={error}")
            self.send_json({"ok": False, "error": str(error)}, status=500)

    def handle_control(self) -> None:
        try:
            length = int(self.headers.get("Content-Length", "0"))
            payload = json.loads(self.rfile.read(length).decode("utf-8"))
            result = control_job(str(payload.get("action") or ""), str(payload.get("jobId") or ""))
            self.send_json(result, status=200 if result.get("ok") else 400)
        except Exception as error:
            log(f"control error={error}")
            self.send_json({"ok": False, "error": str(error)}, status=500)

    def log_message(self, _format: str, *args) -> None:
        return

    def send_json(self, data: dict, status: int = 200) -> None:
        body = json.dumps(data).encode("utf-8")
        self.send_response(status)
        self.send_cors()
        self.send_header("Content-Type", "application/json; charset=utf-8")
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)

    def send_cors(self) -> None:
        self.send_header("Access-Control-Allow-Origin", "*")
        self.send_header("Access-Control-Allow-Methods", "POST, GET, OPTIONS")
        self.send_header("Access-Control-Allow-Headers", "Content-Type")


def main() -> int:
    DOWNLOAD_DIR.mkdir(parents=True, exist_ok=True)
    load_jobs()
    server = ThreadingHTTPServer((HOST, PORT), Handler)
    log(f"start host={HOST} port={PORT} dir={DOWNLOAD_DIR}")
    server.serve_forever()
    return 0


if __name__ == "__main__":
    raise SystemExit(main())
