from __future__ import annotations

import asyncio
import gzip
import json
import os
import re
import sqlite3
import subprocess
import sys
import time
from datetime import datetime, timedelta
from io import BytesIO
from pathlib import Path
from runpy import run_module, run_path
from typing import Any, Dict, Iterator, List, Optional, Sequence, Tuple, Union

import click
import httpx
import msgpack

from . import django_schema
from .config import load_config
from .db import (
    TraceNotFoundError,
    convert_json_to_msgpack,
    create_schema_table,
    db_connection,
    delete_traces_before,
    delete_traces_by_id,
    list_traces_from_db,
    load_trace_from_db,
    save_schema,
    save_trace_in_sqlite,
    setup_db,
    vacuum_db,
)
from .generate_tests import generate_from_trace_ids, generate_test_intermediate_format
from .profiler import enable
from .serialize import load_msgpack, monkeypatch_queryset_repr
from .upload import upload_to_dashboard
from .utils import (
    get_terminal_formatter,
    highlight_python,
    highlight_sql,
    maybe_format,
    pretty_byte_size,
)

DATETIME_FORMATS = click.DateTime(
    (
        "%Y-%m-%d",
        "%Y-%m-%dT%H:%M:%S",
        "%Y-%m-%d %H:%M:%S",
        "%Y-%m-%dT%H:%M:%S.%f",
        "%Y-%m-%d %H:%M:%S.%f",
    )
)


DJANGO_SETTINGS_ERROR = """Django settings not found.
Use `--settings` or set the `DJANGO_SETTINGS_MODULE` environment variable."""

TRACE_NOT_FOUND_ERROR = "Could not find trace_id: `{trace_id}`"


@click.group()
def cli():
    # Ensure the current working directory is on the path.
    # Important when running the `kolo` script installed by setuptools.
    # Not really necessary when using `python -m kolo`, but doesn't hurt.
    # Note, we use 1, not 0: https://stackoverflow.com/q/10095037
    # This probably doesn't matter for our use case, but it doesn't hurt.
    sys.path.insert(1, ".")


def python_noop_profiler(frame, event, arg):  # pragma: no cover
    pass


@cli.command(context_settings={"ignore_unknown_options": True})
@click.argument("path")
@click.argument("args", nargs=-1, type=click.UNPROCESSED)
@click.option(
    "--one-trace-per-test",
    default=False,
    is_flag=True,
    help="Generate a trace for each test traced by Kolo.",
)
@click.option("--noop", default=False, is_flag=True, hidden=True)
@click.option(
    "--use-monitoring",
    default=False,
    is_flag=True,
    help="Enable Kolo implementation using sys.monitoring",
)
def run(path, args, one_trace_per_test, noop, use_monitoring):
    """
    Profile Python code using Kolo.

    PATH is the path to the python file or module being profiled.
    """
    if path == "python":
        path, *args = args
        if path == "-m":
            path, *args = args
            module = True
        else:
            module = False
    elif path.endswith(".py"):
        module = False
    else:
        module = True

    existing_profiler = sys.getprofile()
    if existing_profiler:
        raise click.ClickException(
            f"Cannot run Kolo: {existing_profiler} is already active."
        )

    # Monkeypatch sys.argv, so the profiled code doesn't get confused
    # Without this, the profiled code would see extra args it doesn't
    # know how to handle.
    sys.argv = [path, *args]

    if noop:  # pragma: no cover
        config = load_config()
        if config.get("use_rust", True):
            from ._kolo import register_noop_profiler

            register_noop_profiler()
        else:
            sys.setprofile(python_noop_profiler)

        try:
            if module:
                run_module(path, run_name="__main__", alter_sys=True)
            else:
                run_path(path, run_name="__main__")
        finally:
            sys.setprofile(None)
        return

    if use_monitoring:
        config = {"use_monitoring": True}
    else:
        config = {}
    monkeypatch_queryset_repr()
    with enable(config, source="kolo run", _one_trace_per_test=one_trace_per_test):
        if module:
            run_module(path, run_name="__main__", alter_sys=True)
        else:
            run_path(path, run_name="__main__")


@cli.group()
def trace():
    """
    Subcommands for working with Kolo traces.
    """


@trace.command()
@click.argument("path")
@click.option(
    "--created-at",
    help="Mark this trace as created at this time.",
    type=DATETIME_FORMATS,
)
def load(path, created_at):
    """
    Load a trace from a file into the Kolo database.
    """
    db_path = setup_db()

    try:
        with open(path, "rb") as dump:
            raw_data = dump.read()
    except FileNotFoundError:
        raise click.ClickException(f'File "{path}" not found')

    try:
        data = msgpack.unpackb(raw_data, strict_map_key=False)
    except Exception:
        raise click.ClickException("Trace file is not valid msgpack")

    try:
        trace_id = data["trace_id"]
    except KeyError:
        raise click.ClickException("Trace file is missing the `trace_id`")

    try:
        save_trace_in_sqlite(
            db_path,
            trace_id,
            msgpack=raw_data,
            ignore_errors=False,
            created_at=created_at,
        )
    except sqlite3.IntegrityError:
        raise click.ClickException(f"Trace ID `{trace_id}` already exists")

    click.echo(f"Loaded trace {trace_id}")


@trace.command()
@click.argument("trace_id")
@click.option("--file", help="The name of the file to save the trace to.")
@click.option(
    "--as-python",
    default=False,
    is_flag=True,
    help="Show the trace as readable Python types.",
)
@click.option(
    "--syntax-highlight",
    help="Highlight Python syntax. Ignored unless `--as-python` is set.",
    default="off",
    flag_value="dark",
    is_flag=False,
    type=click.Choice(["off", "light", "dark"], case_sensitive=False),
)
def dump(trace_id, file, as_python, syntax_highlight):
    """
    Dump a trace from the Kolo database to stdout or a specified file.
    """
    db_path = setup_db()

    try:
        msgpack_data, _ = load_trace_from_db(db_path, trace_id)
    except TraceNotFoundError:
        raise click.ClickException(TRACE_NOT_FOUND_ERROR.format(trace_id=trace_id))

    if as_python:
        data = load_msgpack(msgpack_data)
        data = repr(data)
        data = maybe_format(data)
        if syntax_highlight != "off":
            formatter = get_terminal_formatter(syntax_highlight)
            data = highlight_python(data, formatter)
        if file:
            with open(file, "w") as f:
                f.write(data)
        else:
            click.echo(data)
    elif file:
        with open(file, "wb") as f:
            f.write(msgpack_data)

    else:
        click.echo(msgpack_data, nl=False)


@trace.command()
@click.option("--count", help="The number of rows to show.", default=500)
@click.option(
    "--reverse",
    help="Reverse the order of the rows: newest at the bottom.",
    default=False,
    is_flag=True,
)
def list(count, reverse):
    """
    List of recent traces stored by Kolo.
    """
    db_path = setup_db()

    traces = list_traces_from_db(db_path, count=count, reverse=reverse)

    if not traces:
        click.echo("No traces found")
        return

    for trace, timestamp, size in traces:
        size = pretty_byte_size(size)
        click.echo(f"{trace} at {timestamp} ({size})")


@trace.command()
@click.argument("trace_id")
@click.option(
    "--syntax-highlight",
    help="Highlight SQL query syntax.",
    default="off",
    flag_value="dark",
    is_flag=False,
    type=click.Choice(["off", "light", "dark"], case_sensitive=False),
)
def list_queries(trace_id, syntax_highlight):
    """List all SQL queries in a trace."""
    db_path = setup_db()

    try:
        msgpack_data, _ = load_trace_from_db(db_path, trace_id)
    except TraceNotFoundError:
        raise click.ClickException(TRACE_NOT_FOUND_ERROR.format(trace_id=trace_id))

    formatter = get_terminal_formatter(syntax_highlight)

    data = load_msgpack(msgpack_data)
    for frame in data["frames_of_interest"]:
        if frame["type"] == "end_sql_query":
            query = highlight_sql(
                frame["query"], dialect=frame["database"], formatter=formatter
            )
            click.echo(query)


@trace.command()
@click.argument("trace_ids", required=False, nargs=-1)
@click.option("--old", is_flag=True, default=False, help="Delete old traces.")
@click.option(
    "--before",
    help="Delete traces older than this datetime. Must be used with `--old`.",
    type=DATETIME_FORMATS,
)
@click.option(
    "--vacuum",
    help="Recover disk space from the Kolo database.",
    default=False,
    is_flag=True,
)
def delete(trace_ids, old, before, vacuum):
    """
    Delete one or more traces stored by Kolo.
    """

    if before is not None and old is False:
        raise click.ClickException("--before requires --old")

    if old is False and not trace_ids and vacuum is False:
        raise click.ClickException("Must specify either TRACE_IDS, --old, or --vacuum")

    if trace_ids and old:
        raise click.ClickException("Cannot specify TRACE_IDS and --old together")

    db_path = setup_db()

    if trace_ids:
        delete_traces_by_id(db_path, trace_ids)
    elif old:
        if before is None:
            before = datetime.now() - timedelta(days=30)

        deleted_count = delete_traces_before(db_path, before)
        click.echo(f"Deleted {deleted_count} old traces created before {before}.")

    if vacuum:
        vacuum_db(db_path)


def manage_py_settings():
    import ast

    try:
        with open("manage.py") as f:
            data = f.read()
    except OSError:  # pragma: no cover
        return None
    source = ast.parse(data, "manage.py")
    for node in ast.walk(source):
        if (
            isinstance(node, ast.Call)
            and isinstance(node.func, ast.Attribute)
            and node.func.attr == "setdefault"
            and isinstance(node.args[0], ast.Constant)
            and node.args[0].value == "DJANGO_SETTINGS_MODULE"
            and isinstance(node.args[1], ast.Constant)
        ):
            return node.args[1].value  # pragma: no cover
    return None


def load_django(settings):
    import django

    if settings:
        os.environ["DJANGO_SETTINGS_MODULE"] = settings
    elif os.environ.get("DJANGO_SETTINGS_MODULE"):
        pass
    else:
        settings = manage_py_settings()
        if settings:
            os.environ["DJANGO_SETTINGS_MODULE"] = settings  # pragma: no cover
        else:
            raise click.ClickException(DJANGO_SETTINGS_ERROR)

    django.setup()


@cli.command()
@click.argument("trace_ids", nargs=-1)
@click.option(
    "--test-name", default="test_my_view", help="The name of the generated test."
)
@click.option(
    "--test-class",
    default="MyTestCase",
    help="The name of the generated TestCase class.",
)
@click.option(
    "--file",
    type=click.File("w"),
    help="Write the generated test to this file.",
)
@click.option(
    "--settings", default="", help="The dotted path to a Django settings module."
)
@click.option(
    "--use-saved-schemas",
    default=False,
    is_flag=True,
    help="Load Django schemas saved with `kolo store-django-model-schema` instead of using the current schema. This may be useful when generating a test from an old trace.",
)
@click.option(
    "--intermediate-format",
    default=False,
    is_flag=True,
    help="Show an intermediate format of the data Kolo has extracted from the trace for use in the test.",
)
@click.option(
    "--and-run",
    default=False,
    is_flag=True,
    help="[Experimental, Use not recommended] Immediately run the newly generated test.",
)
@click.option(
    "--unittest", "pytest", flag_value=False, help="Generate a unittest style test."
)
@click.option(
    "--pytest",
    "pytest",
    flag_value=True,
    default=True,
    help="Generate a pytest style test (default).",
)
def generate_test(
    trace_ids,
    test_name,
    test_class,
    file,
    settings,
    use_saved_schemas,
    intermediate_format,
    and_run,
    pytest,
):
    """Generate a test from a Kolo trace."""
    import logging

    logging.disable()

    try:
        load_django(settings)

        if intermediate_format:
            for data in generate_test_intermediate_format(
                *trace_ids,
                test_class=test_class,
                test_name=test_name,
                use_saved_schemas=use_saved_schemas,
            ):
                click.echo(data)
            return
        try:
            test_code = generate_from_trace_ids(
                *trace_ids,
                test_class=test_class,
                test_name=test_name,
                use_saved_schemas=use_saved_schemas,
                pytest=pytest,
            )
        except TraceNotFoundError as e:
            raise click.ClickException(TRACE_NOT_FOUND_ERROR.format(trace_id=e.args[0]))

        if and_run:
            import unittest

            from django.test.runner import DiscoverRunner

            our_globals: Dict[Any, Any] = {}
            exec(test_code, our_globals)

            # Create TestSuite manually and add test cases
            test_suite = unittest.TestSuite()
            test_suite.addTest(
                unittest.defaultTestLoader.loadTestsFromTestCase(
                    our_globals[test_class]
                )
            )

            class CustomRunner(DiscoverRunner):
                def build_suite(
                    self, test_labels: Optional[Sequence[str]] = None, **kwargs
                ) -> unittest.TestSuite:
                    return test_suite

                def suite_result(self, suite, result, **kwargs):
                    return suite, result

            runner = CustomRunner()

            # Test output will be output in the terminal as usual
            suite, result = runner.run_tests([])  # type: ignore # we're now returning a tuple (see CustomRunner.suite_result above)
        elif file:
            file.write(test_code)
        else:
            click.echo(test_code)
    finally:
        logging.disable(logging.NOTSET)


@cli.command()
@click.option(
    "--settings", default="", help="The dotted path to a Django settings module."
)
def store_django_model_schema(settings):
    """Store Django model info for test generation."""
    from .git import COMMIT_SHA

    load_django(settings)

    schema = django_schema.get_schema()

    db_path = setup_db()
    with db_connection(db_path) as connection:
        create_schema_table(connection)
        save_schema(connection, schema, COMMIT_SHA)


@cli.command()
def dbshell():  # pragma: no cover
    """
    Open a sqlite3 shell to the Kolo database.
    """
    db_path = setup_db()
    subprocess.run(["sqlite3", db_path], check=True)


@trace.command()
@click.argument("trace_id")
def upload(trace_id):
    """
    Upload a trace to the Kolo dashboard
    """
    db_path = setup_db()

    try:
        msgpack_data, _ = load_trace_from_db(db_path, trace_id)
    except TraceNotFoundError:
        raise click.ClickException(TRACE_NOT_FOUND_ERROR.format(trace_id=trace_id))

    response = upload_to_dashboard(msgpack_data)

    if response.status_code == 201:
        click.echo(f"{trace_id} uploaded successfully!")
    else:
        errors = response.json()
        raise click.ClickException(errors)


@trace.command()
@click.argument("trace_id")
def download(trace_id):
    """
    Download a trace from the Kolo dashboard
    """
    db_path = setup_db()

    # TODO(later): The ability to download will ultimately likely be guarded by authenticating
    # against a given project/organisation etc so that you cannot download someone else's trace
    # just by guessing a ULID, so we'll probably need to pass those along eventually too...
    base_url = os.environ.get("KOLO_BASE_URL", "https://my.kolo.app")
    url = f"{base_url}/api/traces/{trace_id}/download"
    response = httpx.get(url)

    if response.status_code == 404:
        raise click.ClickException(f"`{trace_id}` was not found by the server.")
    elif response.status_code != 200:
        raise click.ClickException(f"Unexpected status code: {response.status_code}.")

    raw_data = response.content

    try:
        msgpack.unpackb(raw_data, strict_map_key=False)
    except Exception:
        raise click.ClickException("Downloaded trace was not valid msgpack.")

    try:
        save_trace_in_sqlite(
            db_path,
            trace_id,
            msgpack=raw_data,
            ignore_errors=False,
        )
    except sqlite3.IntegrityError:
        raise click.ClickException(f"`{trace_id}` already exists.")

    click.echo(f"`{trace_id}` downloaded successfully!")


@trace.command()
def json_to_msgpack():  # pragma: no cover
    """
    Convert all legacy json traces to msgpack
    """
    db_path = setup_db()
    count = convert_json_to_msgpack(db_path)
    click.echo(f"{count} traces converted!")


@cli.group()
def ci():
    """
    Subcommands for CI-related operations.
    """
    pass


def compress_trace(trace_data: bytes) -> bytes:
    compressed_data = BytesIO()
    with gzip.GzipFile(fileobj=compressed_data, mode="wb") as gz:
        gz.write(trace_data)
    return compressed_data.getvalue()


def create_trace_group() -> str:
    auth_token = os.environ["KOLO_TOKEN"]
    headers = {"Authorization": f"Bearer {auth_token}"}
    base_url = os.environ.get("KOLO_BASE_URL", "https://my.kolo.app")
    response = httpx.post(
        f"{base_url}/api/trace-groups", headers=headers, json={"name": "ci"}
    )
    if response.status_code != 201:
        error_message = f"Failed to create trace group. Status code: {response.status_code}\n{response.text}"

        click.echo(click.style(error_message, fg="red"), err=True)
        raise click.Abort()

    trace_group = response.json()

    return trace_group["id"]


async def upload_trace(
    client: Union[httpx.Client, httpx.AsyncClient],
    trace_id: str,
    upload_url: str,
    db_path: Path,
    max_retries: int = 3,
    initial_backoff: float = 1.0,
) -> Optional[str]:
    loaded_trace = load_trace_from_db(db_path, trace_id)

    msgpack = loaded_trace[0]
    compressed_data = compress_trace(msgpack)

    for attempt in range(max_retries):
        try:
            if isinstance(client, httpx.AsyncClient):
                response = await client.put(
                    upload_url,
                    content=compressed_data,
                    headers={"Content-Type": "application/gzip"},
                    timeout=30.0,
                )
            else:
                response = client.put(
                    upload_url,
                    content=compressed_data,
                    headers={"Content-Type": "application/gzip"},
                    timeout=30.0,
                )

            if response.status_code != 200:
                error_message = (
                    f"Failed to upload {trace_id}.msgpack.gz. Status code: {response.status_code}\n"
                    f"{response.text}"
                )
                click.echo(click.style(error_message, fg="red"), err=True)
                return error_message

            return None  # Success case

        except (httpx.TimeoutException, httpx.RequestError) as e:
            if attempt < max_retries - 1:
                backoff = initial_backoff * (2**attempt)
                error_message = f"Error uploading {trace_id}.msgpack.gz (attempt {attempt + 1}/{max_retries}): {str(e)}. Retrying in {backoff:.2f} seconds..."
                click.echo(click.style(error_message, fg="yellow"), err=True)
                if isinstance(client, httpx.AsyncClient):
                    await asyncio.sleep(backoff)
                else:
                    time.sleep(backoff)
            else:
                error_message = f"Failed to upload {trace_id}.msgpack.gz after {max_retries} attempts: {str(e)}"
                click.echo(click.style(error_message, fg="red"), err=True)
                return error_message

    return f"Failed to upload {trace_id}.msgpack.gz after {max_retries} attempts"


def sync_ci_upload(traces: List[Dict[str, Any]], db_path: Path):
    with httpx.Client() as client:
        with click.progressbar(traces, label="Uploading traces", show_pos=True) as bar:
            for trace in bar:
                result = asyncio.run(
                    upload_trace(
                        client=client,
                        trace_id=trace["id"],
                        upload_url=trace["upload_url"],
                        db_path=db_path,
                    )
                )
                if result:
                    click.echo(f"\n{result}", err=True)


async def async_ci_upload(traces: List[Dict[str, Any]], db_path: Path):
    async with httpx.AsyncClient(limits=httpx.Limits(max_connections=20)) as client:
        tasks = [
            upload_trace(
                client=client,
                trace_id=trace["id"],
                upload_url=trace["upload_url"],
                db_path=db_path,
            )
            for trace in traces
        ]

        bar: click._termui_impl.ProgressBar[Any]
        with click.progressbar(
            length=len(tasks), label="Uploading traces", show_pos=True
        ) as bar:
            results = []
            for task in asyncio.as_completed(tasks):
                result = await task
                bar.update(1)
                if result:
                    results.append(result)
                    click.echo(f"\n{result}", err=True)


@ci.command("upload")
@click.option("--sync", is_flag=True, help="Use synchronous upload instead of async")
def ci_upload(sync: bool):
    """
    Upload all traces in the local Kolo db to Kolo Cloud
    """
    db_path = setup_db()
    auth_token = os.environ["KOLO_TOKEN"]

    traces_in_local_db = list_traces_from_db(db_path, count=10000)

    trace_group_id = create_trace_group()

    traces_registered = []
    page_size = 500

    for i in range(0, len(traces_in_local_db), page_size):
        page = traces_in_local_db[i : i + page_size]

        base_url = os.environ.get("KOLO_BASE_URL", "https://my.kolo.app")
        response = httpx.post(
            f"{base_url}/api/traces/bulk",
            headers={"Authorization": f"Bearer {auth_token}"},
            json={
                "trace_group_id": trace_group_id,
                "traces": [{"id": trace[0], "recorded_at": trace[1]} for trace in page],
            },
        )

        if response.status_code != 201:
            error_message = f"Failed to bulk create traces. Status code: {response.status_code}\n{response.text}"
            click.echo(click.style(error_message, fg="red"), err=True)
            raise click.Abort()

        traces_registered.extend(response.json()["traces"])

        click.echo(
            f"Registered {min(i + page_size, len(traces_in_local_db))} of {len(traces_in_local_db)} traces"
        )

    if sync:
        sync_ci_upload(traces=traces_registered, db_path=db_path)
    else:
        asyncio.run(async_ci_upload(traces=traces_registered, db_path=db_path))


if __name__ == "__main__":
    cli()  # pragma: no cover
