import random
import time
from copy import deepcopy
from threading import Lock

import requests

from ai_testing_swarm.core.config import (
    AI_SWARM_RETRY_BACKOFF_MS,
    AI_SWARM_RETRY_COUNT,
    AI_SWARM_RPS,
)


class ExecutionAgent:
    # Shared limiter across threads in a run to enforce true global RPS.
    _throttle_lock = Lock()
    _next_allowed_ts = 0.0

    def _set_path(self, obj: dict, path: list[str], value, remove: bool = False):
        """Set/remove nested key in a dict. Creates intermediate dicts if needed."""
        if not path:
            return
        cur = obj
        for key in path[:-1]:
            if key not in cur or not isinstance(cur[key], dict):
                cur[key] = {}
            cur = cur[key]

        last = path[-1]
        if remove:
            cur.pop(last, None)
        else:
            cur[last] = value

    def apply_mutation(self, data: dict, mutation: dict):
        if not mutation:
            return deepcopy(data)

        data = deepcopy(data)
        path = mutation.get("path") or []
        op = mutation.get("operation")

        if not path:
            if op == "REMOVE":
                return {}
            if op == "REPLACE":
                return deepcopy(mutation.get("new_value"))
            return data

        if op == "REMOVE":
            self._set_path(data, path, None, remove=True)
        elif op == "REPLACE":
            self._set_path(data, path, mutation.get("new_value"), remove=False)

        return data

    def _throttle(self):
        # Optional global throttling to avoid hammering target services.
        # This is process-level and shared by all worker threads.
        if AI_SWARM_RPS and AI_SWARM_RPS > 0:
            interval = 1.0 / float(AI_SWARM_RPS)
            cls = type(self)
            with cls._throttle_lock:
                now = time.monotonic()
                sleep_for = cls._next_allowed_ts - now
                if sleep_for > 0:
                    time.sleep(sleep_for)
                    now = time.monotonic()
                cls._next_allowed_ts = max(now, cls._next_allowed_ts) + interval

    def _should_retry(self, status_code: int | None, error: str | None) -> bool:
        if error:
            return True
        if status_code in (408, 425, 429, 500, 502, 503, 504):
            return True
        return False

    def execute(self, request: dict, test: dict):
        mutation = test.get("mutation")

        # Method (allow method mutations)
        method = request["method"].upper()
        if mutation and mutation.get("target") == "method" and mutation.get("operation") == "REPLACE":
            method = str(mutation.get("new_value") or method).upper()

        params = (
            self.apply_mutation(request.get("params", {}) or {}, mutation)
            if mutation and mutation.get("target") == "query"
            else deepcopy(request.get("params", {}) or {})
        )

        body = deepcopy(request.get("body", {}) or {})
        if mutation and mutation.get("target") == "body":
            body = self.apply_mutation(body, mutation)
        elif mutation and mutation.get("target") == "body_whole" and mutation.get("operation") == "REPLACE":
            body = mutation.get("new_value")

        headers = deepcopy(request.get("headers", {}) or {})
        if mutation and mutation.get("target") == "headers":
            headers = self.apply_mutation(headers, mutation)

        # Payload strategy (basic content-type aware)
        send_json = method != "GET"
        ct = (headers.get("content-type") or headers.get("Content-Type") or "").lower()

        json_payload = None
        data_payload = None
        if method == "GET":
            json_payload = None
            data_payload = None
        else:
            # If body is not a dict, send raw data
            if not isinstance(body, (dict, list)):
                send_json = False
                data_payload = body
            elif "application/x-www-form-urlencoded" in ct:
                send_json = False
                data_payload = body
            else:
                json_payload = body

        # Execute with light retry/backoff on transient failures.
        attempt = 0

        while True:
            attempt += 1
            self._throttle()

            started = time.time()
            try:
                resp = requests.request(
                    method=method,
                    url=request["url"],
                    headers=headers,
                    params=params or None,
                    json=json_payload if send_json else None,
                    data=data_payload if not send_json else None,
                    timeout=15,
                )
                elapsed_ms = int((time.time() - started) * 1000)

                if attempt <= (AI_SWARM_RETRY_COUNT + 1) and self._should_retry(resp.status_code, None):
                    # jittered backoff
                    backoff = (AI_SWARM_RETRY_BACKOFF_MS / 1000.0) * (2 ** (attempt - 1))
                    backoff = backoff * (0.75 + random.random() * 0.5)
                    time.sleep(min(backoff, 5.0))
                    if attempt <= AI_SWARM_RETRY_COUNT:
                        continue

                try:
                    body_snippet = resp.json()
                except Exception:
                    body_snippet = resp.text
                response_headers = dict(getattr(resp, "headers", {}) or {})

                openapi_validation = []
                try:
                    ctx = request.get("_openapi") or {}
                    if isinstance(ctx, dict) and ctx.get("spec") and ctx.get("path") and ctx.get("method"):
                        # If method was mutated to something else, we can't reliably map to an OpenAPI op.
                        if str(ctx.get("method")).upper() == str(method).upper():
                            from ai_testing_swarm.core.openapi_validator import validate_openapi_response

                            issues = validate_openapi_response(
                                spec=ctx["spec"],
                                path=ctx["path"],
                                method=ctx["method"],
                                status_code=resp.status_code,
                                response_headers=response_headers,
                                response_json=body_snippet if isinstance(body_snippet, (dict, list)) else None,
                            )
                            openapi_validation = [i.__dict__ for i in issues]
                except Exception:
                    # Best-effort only; never fail execution on validator issues.
                    openapi_validation = []

                return {
                    "name": test["name"],
                    "mutation": mutation,
                    "request": {
                        "method": method,
                        "url": request["url"],
                        "headers": headers,
                        "params": params,
                        "body": body if method != "GET" else None,
                    },
                    "response": {
                        "status_code": resp.status_code,
                        "elapsed_ms": elapsed_ms,
                        "attempt": attempt,
                        "headers": response_headers,
                        "body_snippet": body_snippet,
                        "openapi_validation": openapi_validation,
                    },
                }

            except requests.RequestException as e:
                elapsed_ms = int((time.time() - started) * 1000)
                last_error = str(e)

                if attempt <= AI_SWARM_RETRY_COUNT and self._should_retry(None, last_error):
                    backoff = (AI_SWARM_RETRY_BACKOFF_MS / 1000.0) * (2 ** (attempt - 1))
                    backoff = backoff * (0.75 + random.random() * 0.5)
                    time.sleep(min(backoff, 5.0))
                    continue

                return {
                    "name": test["name"],
                    "mutation": mutation,
                    "request": {
                        "method": method,
                        "url": request["url"],
                        "headers": headers,
                        "params": params,
                        "body": body if method != "GET" else None,
                    },
                    "response": {
                        "status_code": None,
                        "elapsed_ms": elapsed_ms,
                        "attempt": attempt,
                        "error": last_error,
                        "body_snippet": None,
                    },
                }
