import argparse
import math
import os
import json
import sys
import requests
import hashlib
import base64

from pycafe_server.client import requirements_system, shared

app_types = ["vizro", "dash", "streamlit", "solara", "panel", "shiny"]
pyodide_versions = ["0.27.2", "0.27.1", "0.26.2", "0.26.1", "0.25.1", "0.23.0"]
public_url = os.environ.get("PYCAFE_CLIENT_PUBLIC_URL") or "https://py.cafe"


def main():
    parser = argparse.ArgumentParser(
        description="Export a directory with an app to a standalone HTML file."
    )
    parse_export_html(parser)
    args = parser.parse_args()
    run_export_html(args)


def parse_export_html(parser: argparse.ArgumentParser):
    parser.add_argument(
        "--input-dir",
        default="",
        help="The source directory to make html (default: '')",
        required=True,
    )
    parser.add_argument(
        "--name", default="project", help="The of the html (default: 'project')"
    )
    parser.add_argument("--title", default="", help="The title (default: '')")
    parser.add_argument(
        "--pyodide",
        default="0.26.1",
        choices=pyodide_versions,
        help=f"The pyodide version (default: '0.26.1') {pyodide_versions}",
    )
    parser.add_argument(
        "--type",
        default="",
        choices=app_types,
        help=f"The app type. (default: '') {app_types}",
        required=True,
    )
    parser.add_argument(
        "--public",
        action="store_true",
        help="Export using public py.cafe as viewer, unsigned (default: False)",
    )


def run_export_html(args: argparse.Namespace):
    url = (
        public_url
        if args.public
        else os.environ.get("PYCAFE_CLIENT_URL", "").rstrip("/")
    )
    token = "" if args.public else os.environ.get("PYCAFE_API_KEY", "")
    if not args.public and (not url or not token):
        print(
            "Error: PYCAFE_CLIENT_URL and PYCAFE_API_KEY environment variables are required when not using --public mode"
        )
        sys.exit(1)

    if args.public:
        print("Exporting using public py.cafe viewer (unsigned).")

    module_dir = os.path.dirname(__file__)
    file_path = os.path.join(module_dir, "file-template.html")

    with open(file_path, "r") as f:
        template = f.read()

    project = {
        "name": args.name,
        "type": args.type,
        "python": f"pyodide-v{args.pyodide}",
        "code": None,
        "requirements": None,
        "files": [],
        "title": args.title,
    }

    html_export_version = 1
    try:
        resp = requests.get(f"{url}/api2/versions", timeout=3)
        resp.raise_for_status()
        versions = resp.json()
        if "htmlExport" in versions:
            html_export_version = int(versions["htmlExport"])
    except Exception as e:
        print(
            f"Could not get versions from server, defaulting to HTML export version {html_export_version}.  ({e})"
        )

    MAX_PROJECT_SIZE = (100 if html_export_version == 1 else 400) * 1000 * 1000
    SIZE_THRESHOLD = 32 * 1024
    total_size = 0

    for root, dirs, files in os.walk(args.input_dir):
        dirs[:] = [
            d
            for d in dirs
            if d not in ["venv", ".venv", "env", ".env", ".git", "__pycache__"]
        ]
        for filename in files:
            file_path = os.path.join(root, filename)
            sub_path = file_path[len(args.input_dir) + 1 :]
            try:
                if filename in [".DS_Store"]:
                    continue
                if sub_path in ["app.py", "requirements.txt"]:
                    with open(file_path, "r", encoding="utf-8") as f:
                        project["code" if sub_path == "app.py" else "requirements"] = (
                            f.read()
                        )
                else:
                    encoding = {}
                    base64_encode = False
                    total_size += os.path.getsize(file_path)

                    # we can only use large file chunking with base64 encoding from v2
                    if (
                        os.path.getsize(file_path) < SIZE_THRESHOLD
                        and html_export_version >= 2
                    ):
                        try:
                            with open(file_path, "r", encoding="utf-8") as f:
                                content = f.read()
                        except UnicodeDecodeError:
                            base64_encode = True
                    else:
                        base64_encode = True

                    if base64_encode:
                        with open(file_path, "rb") as f:
                            content = f.read()
                            content = base64.b64encode(content).decode("utf-8")
                            encoding = {"encoding": "base64"}
                    project["files"].append(
                        {"name": sub_path, "content": content, **encoding}
                    )
            except Exception as e:
                print(f"Could not read {file_path}: {e}")

    def display_size(size):
        return f"{math.ceil(size / 1000 / 1000)}MB"

    if total_size > MAX_PROJECT_SIZE:
        print(
            f"Error: Total project size {display_size(total_size)} exceeds current maximum supported size of {display_size(MAX_PROJECT_SIZE)}"
        )
        if html_export_version > 1:
            exit(1)

    if project["code"] is None or project["requirements"] is None:
        raise ValueError("app.py and requirements.txt are required")

    requirements, _ = shared.parse_requirements(project["requirements"])
    defaults = shared.parse_requirements(getattr(requirements_system, args.type))[0]
    requirements = shared.add_requirements_default(requirements, defaults)

    # lock_file = "https://cdn.jsdelivr.net/npm/pyodide@0.27.1/pyodide-lock.json"
    # read lock_file
    lock_file = requests.get(
        f"https://cdn.jsdelivr.net/npm/pyodide@{args.pyodide}/pyodide-lock.json"
    ).json()
    python_version = lock_file["info"]["python"]
    repodata_packages = lock_file["packages"]
    # print(lock_file)

    query, unvendored_but_requested, local_files = shared.create_query(
        requirements, repodata_packages, python_version, url
    )

    files = [
        (
            "wheels",
            (
                file_name[7:],
                shared.create_minimal_wheel(args.input_dir + "/" + file_name[7:]),
            ),
        )
        for file_name in local_files
    ]

    resolve_request = requests.post(
        f"{url}/api/resolve",
        files=files,
        data={"query": query},
    )
    if not resolve_request.ok:
        print(
            f"Error resolving requirements: {resolve_request.status_code} {resolve_request.text}"
        )
        sys.exit(1)
    env_lock = resolve_request.text
    project["lockfileContent"] = env_lock

    sorted_project = dict(sorted(project.items()))
    sorted_project["files"] = [
        dict(sorted(file.items())) for file in sorted_project["files"]
    ]
    project_string = json.dumps(
        sorted_project, separators=(",", ":"), ensure_ascii=False
    )
    hash = hashlib.sha256(project_string.encode()).hexdigest()

    try:
        if args.public:
            signature = None
        else:
            response = requests.get(
                f"{url}/api/sign?hash={hash}",
                headers={
                    "Authorization": f"Bearer {token}",
                    "X-Authorization": f"Bearer {token}",
                },
            )
            if not response.ok:
                print(f"Error: {response.status_code} {response.text}")
                sys.exit(1)
            signature = response.json()["signatureJwt"]

        def should_chunk(file):
            return (
                html_export_version >= 2
                and "encoding" in file
                and file["encoding"] == "base64"
                and len(file["content"]) > SIZE_THRESHOLD
            )

        def remove_content(file):
            if should_chunk(file):
                return {k: v for k, v in file.items() if k != "content"}
            return file

        files_to_chunk = [
            file for file in sorted_project["files"] if should_chunk(file)
        ]

        sorted_project["files"] = [
            remove_content(file) for file in sorted_project["files"]
        ]
        project_string = json.dumps(
            sorted_project, separators=(",", ":"), ensure_ascii=False
        )

        with open(f"{args.name}.html", "w") as f:
            f.write(
                template.replace("#title#", args.name)
                .replace("#viewerUrl#", f"{url}/view?iframe")
                .replace('"#signatureJwt#"', f'"{signature}"' if signature else "null")
                .replace("true; //#preAuth#", "false")
                .replace("#origin#", url)
                .replace('{ "projectState": "" }', project_string.replace("/", "\\/"))
                .replace(
                    "<!-- FILE_PARTS -->",
                    "\n".join(
                        [
                            make_chunks(file["name"], file["content"])
                            for file in files_to_chunk
                        ]
                    ),
                )
            )
        print(f"Created {args.name}.html")
    except requests.RequestException as e:
        print("Error occurred while making the request:", e)


def make_chunks(name: str, base64_str: str) -> str:
    CHUNK_SIZE = 2 * 1024 * 1024

    if len(base64_str) % 4 != 0:
        print(
            "make_chunks: input length is not a multiple of 4; chunks may be invalid Base64."
        )

    out_parts = []
    i = 0
    n = len(base64_str)
    safe_name = base64.b64encode(name.encode("utf-8")).decode("utf-8")

    while i < n:
        # Take up to CHUNK_SIZE chars, then trim to a multiple of 4
        length = min(CHUNK_SIZE, n - i)
        length -= length % 4

        # If we're at the end and length became 0 due to modulo,
        # just take the remaining (should only happen if remaining < 4).
        if length == 0:
            length = n - i

        chunk = base64_str[i : i + length]
        out_parts.append(
            f'<script type="application/octet-stream" data-filename="{safe_name}">{chunk}</script>\n'
        )
        i += length

    return "".join(out_parts)


if __name__ == "__main__":
    main()
