"""K6 rules K6001-K6006."""

from __future__ import annotations

import re

from perf_lint.ir.models import Framework, Location, Severity, Violation
from perf_lint.ir.models import ScriptIR
from perf_lint.rules.base import BaseRule, RuleRegistry


@RuleRegistry.register
class K6001MissingThinkTime(BaseRule):
    rule_id = "K6001"
    name = "MissingThinkTime"
    description = "No sleep() calls found. Without think time, virtual users hammer the server as fast as possible, creating unrealistic load."
    severity = Severity.WARNING
    frameworks = [Framework.K6]
    tags = ("think-time", "realism")
    fixable = True

    def check(self, ir: ScriptIR) -> list[Violation]:
        if not ir.parsed_data.get("has_sleep"):
            return [
                Violation(
                    rule_id=self.rule_id,
                    severity=self.severity,
                    message="No sleep() calls found. Add think time to simulate realistic user behaviour.",
                    location=Location(line=1),
                    suggestion="Import sleep from 'k6' and add sleep(1) between requests to simulate user think time.",
                    fix_example="import { sleep } from 'k6';\n\nexport default function () {\n  http.get('https://example.com');\n  sleep(1); // 1 second think time\n}",
                )
            ]
        return []

    def apply_fix(self, ir: ScriptIR) -> str | None:
        if ir.parsed_data.get("has_sleep"):
            return None  # Already has sleep — nothing to fix
        source = ir.raw_content

        # 1. Ensure sleep is imported from 'k6'
        k6_import_re = re.compile(r"(import\s*\{)([^}]*?)(\}\s*from\s*['\"]k6['\"])")
        k6_match = k6_import_re.search(source)
        if k6_match:
            imports_str = k6_match.group(2)
            if "sleep" not in imports_str:
                # Add sleep to existing import
                new_imports = imports_str.rstrip() + ", sleep"
                source = (
                    source[: k6_match.start(2)]
                    + new_imports
                    + source[k6_match.end(2) :]
                )
        else:
            # No k6 import found; add one at the top
            source = "import { sleep } from 'k6';\n" + source

        # 2. Insert sleep(1) before the closing brace of export default function
        func_re = re.compile(r"export\s+default\s+function\s*\w*\s*\([^)]*\)\s*\{")
        func_match = func_re.search(source)
        if not func_match:
            return None

        # NOTE: The brace-depth counter does not account for '{' or '}' characters
        # inside string literals, template literals, or comments. For the vast
        # majority of real K6 scripts this is safe; pathological cases may result
        # in sleep(1) being inserted at the wrong position.

        # Walk forward tracking brace depth to find the matching closing }
        start = func_match.end()
        depth = 1
        pos = start
        while pos < len(source) and depth > 0:
            ch = source[pos]
            if ch == "{":
                depth += 1
            elif ch == "}":
                depth -= 1
            pos += 1

        if depth != 0:
            return None  # Unmatched braces — bail out

        closing_pos = pos - 1
        # Determine indentation of the closing brace line
        line_start = source.rfind("\n", 0, closing_pos) + 1
        indent = len(source[line_start:closing_pos]) - len(
            source[line_start:closing_pos].lstrip()
        )
        pad = " " * (indent + 2)
        source = source[:closing_pos] + f"{pad}sleep(1);\n" + source[closing_pos:]
        return source


@RuleRegistry.register
class K6002HardcodedURL(BaseRule):
    rule_id = "K6002"
    name = "HardcodedURL"
    description = "HTTP call uses a hardcoded IP address. This prevents running the script against different environments."
    severity = Severity.ERROR
    frameworks = [Framework.K6]
    tags = ("parameterization", "portability")
    tier = "pro"

    def check(self, ir: ScriptIR) -> list[Violation]:
        violations = []
        for call in ir.parsed_data.get("http_calls", []):
            if call.get("is_hardcoded_ip"):
                violations.append(
                    Violation(
                        rule_id=self.rule_id,
                        severity=self.severity,
                        message=f"HTTP {call['method']} uses hardcoded IP URL: {call['url']!r}. Use __ENV or a configurable base URL.",
                        location=Location(line=None),
                        suggestion="Use __ENV.BASE_URL or a top-level constant to define the base URL.",
                        fix_example="const BASE_URL = __ENV.BASE_URL || 'https://staging.example.com';\n\nhttp.get(`${BASE_URL}/api/endpoint`);",
                    )
                )
        return violations


@RuleRegistry.register
class K6003MissingCheck(BaseRule):
    rule_id = "K6003"
    name = "MissingCheck"
    description = "No check() calls found. Without checks, K6 cannot detect application errors — all requests appear successful."
    severity = Severity.ERROR
    frameworks = [Framework.K6]
    tags = ("assertions", "correctness")
    fixable = True
    tier = "pro"

    def check(self, ir: ScriptIR) -> list[Violation]:
        if not ir.parsed_data.get("has_check"):
            return [
                Violation(
                    rule_id=self.rule_id,
                    severity=self.severity,
                    message="No check() calls found. Tests that can't detect failures aren't tests.",
                    location=Location(line=1),
                    suggestion="Use check() to validate HTTP status codes and response content.",
                    fix_example="import { check } from 'k6';\n\nconst res = http.get('https://example.com');\ncheck(res, {\n  'status is 200': (r) => r.status === 200,\n  'response time < 500ms': (r) => r.timings.duration < 500,\n});",
                )
            ]
        return []

    def apply_fix(self, ir: ScriptIR) -> str | None:
        source = ir.raw_content

        # 1. Ensure check is imported from 'k6'
        k6_import_re = re.compile(r"(import\s*\{)([^}]*?)(\}\s*from\s*['\"]k6['\"])")
        k6_match = k6_import_re.search(source)
        if k6_match:
            imports_str = k6_match.group(2)
            if "check" not in imports_str:
                new_imports = imports_str.rstrip() + ", check"
                source = source[:k6_match.start(2)] + new_imports + source[k6_match.end(2):]
        else:
            source = "import { check } from 'k6';\n" + source

        # 2. Find first http.method( call assigned to a variable
        assign_re = re.compile(
            r"(const|let|var)\s+(\w+)\s*=\s*http\.(get|post|put|del|patch|head|options)\s*\(",
            re.MULTILINE,
        )
        assign_match = assign_re.search(source)
        if assign_match:
            var_name = assign_match.group(2)
            stmt_end = source.find("\n", assign_match.start())
            if stmt_end == -1:
                stmt_end = len(source)
            line_start = source.rfind("\n", 0, assign_match.start()) + 1
            indent = assign_match.start() - line_start
            pad = " " * indent
            check_line = (
                f"\n{pad}check({var_name}, {{\n"
                f"{pad}  'status is 200': (r) => r.status === 200,\n"
                f"{pad}}});"
            )
            source = source[:stmt_end] + check_line + source[stmt_end:]
            return source

        # NOTE: This fix only wraps single-line HTTP calls. Multi-line calls
        # (e.g., http.post(url,\n  body)) are left unchanged to avoid producing
        # syntactically invalid output.

        # If no assignment found, wrap first bare http call
        bare_re = re.compile(
            r"http\.(get|post|put|del|patch|head|options)\s*\(",
            re.MULTILINE,
        )
        bare_match = bare_re.search(source)
        if not bare_match:
            return None
        line_start = source.rfind("\n", 0, bare_match.start()) + 1
        indent = bare_match.start() - line_start
        pad = " " * indent
        stmt_end = source.find("\n", bare_match.start())
        if stmt_end == -1:
            stmt_end = len(source)
        original_line = source[line_start:stmt_end].lstrip()
        source = (
            source[:line_start]
            + f"{pad}const _res = {original_line}\n"
            + f"{pad}check(_res, {{ 'status is 200': (r) => r.status === 200 }});"
            + source[stmt_end:]
        )
        return source


@RuleRegistry.register
class K6004MissingThresholds(BaseRule):
    rule_id = "K6004"
    name = "MissingThresholds"
    description = "No thresholds defined in options. Without thresholds, K6 won't fail the test when SLOs are breached."
    severity = Severity.WARNING
    frameworks = [Framework.K6]
    tags = ("thresholds", "slo", "ci-integration")
    fixable = True

    _THRESHOLDS_BLOCK = (
        "  thresholds: {\n"
        "    http_req_duration: ['p(95)<500'],\n"
        "    http_req_failed: ['rate<0.01'],\n"
        "  },\n"
    )

    def check(self, ir: ScriptIR) -> list[Violation]:
        if not ir.parsed_data.get("has_thresholds"):
            return [
                Violation(
                    rule_id=self.rule_id,
                    severity=self.severity,
                    message="No thresholds defined. Add thresholds to define pass/fail criteria and enforce SLOs in CI.",
                    location=Location(line=1),
                    suggestion="Define thresholds in the options object to set performance SLOs.",
                    fix_example="export const options = {\n  thresholds: {\n    http_req_duration: ['p(95)<500'], // 95% of requests must complete < 500ms\n    http_req_failed: ['rate<0.01'],    // Error rate must be < 1%\n  },\n};",
                )
            ]
        return []

    def apply_fix(self, ir: ScriptIR) -> str | None:
        source = ir.raw_content

        # Try to inject into an existing options object
        options_re = re.compile(r"export\s+const\s+options\s*=\s*\{")
        options_match = options_re.search(source)
        if options_match:
            insert_pos = options_match.end()
            source = source[:insert_pos] + "\n" + self._THRESHOLDS_BLOCK + source[insert_pos:]
        else:
            # No options block — prepend a minimal one
            options_block = (
                "export const options = {\n"
                + self._THRESHOLDS_BLOCK
                + "};\n\n"
            )
            source = options_block + source

        return source


@RuleRegistry.register
class K6005MissingErrorHandling(BaseRule):
    rule_id = "K6005"
    name = "MissingErrorHandling"
    description = "No error handling found. Without error handling, failed requests may cause unexpected script behaviour."
    severity = Severity.WARNING
    frameworks = [Framework.K6]
    tags = ("error-handling", "resilience")
    tier = "pro"

    def check(self, ir: ScriptIR) -> list[Violation]:
        if not ir.parsed_data.get("has_error_handling"):
            return [
                Violation(
                    rule_id=self.rule_id,
                    severity=self.severity,
                    message="No error handling detected. Add checks or try/catch blocks to handle request failures gracefully.",
                    location=Location(line=1),
                    suggestion="Check response status codes and handle errors to make your script resilient.",
                    fix_example="const res = http.get(url);\nif (res.status !== 200) {\n  console.error(`Request failed: ${res.status}`);\n}",
                )
            ]
        return []


@RuleRegistry.register
class K6006AggressiveStages(BaseRule):
    rule_id = "K6006"
    name = "AggressiveStages"
    description = "First stage ramps up too many users too quickly (> 100 users in < 10 seconds), creating an unrealistic spike."
    severity = Severity.WARNING
    frameworks = [Framework.K6]
    tags = ("rampup", "realism")
    fixable = True
    tier = "team"

    def check(self, ir: ScriptIR) -> list[Violation]:
        stages = ir.parsed_data.get("stages", [])
        if not stages:
            return []

        first_stage = stages[0]
        target = first_stage.get("target", 0)
        duration_secs = first_stage.get("duration_secs", 60)

        if target > 100 and duration_secs < 10:
            return [
                Violation(
                    rule_id=self.rule_id,
                    severity=self.severity,
                    message=(
                        f"First stage ramps to {target} users in {duration_secs}s. "
                        "This is an aggressive spike — consider a longer warm-up."
                    ),
                    location=Location(line=None),
                    suggestion=f"Extend the first stage duration to at least {target // 10}s to ramp {target} users gradually.",
                    fix_example=f"stages: [\n  {{ duration: '{target // 10}s', target: {target} }},  // Gradual ramp-up\n  {{ duration: '5m', target: {target} }},              // Steady state\n  {{ duration: '30s', target: 0 }},                   // Ramp-down\n],",
                )
            ]
        return []

    def apply_fix(self, ir: ScriptIR) -> str | None:
        stages = ir.parsed_data.get("stages", [])
        if not stages:
            return None
        first = stages[0]
        target = first.get("target", 0)
        duration_secs = first.get("duration_secs", 60)
        if not (target > 100 and duration_secs < 10):
            return None

        new_duration = max(target // 10, 30)  # At least 30 seconds
        source = ir.raw_content

        # Match the first stage pattern: { duration: '...', target: N }
        _qt = r"""['"]\s*,\s*target\s*:\s*"""
        pattern = r"\{\s*duration\s*:\s*['\"](\d+)([smh])" + _qt + "(" + str(target) + r")\s*\}"
        stage_re = re.compile(pattern, re.MULTILINE)
        m = stage_re.search(source)
        if not m:
            return None

        new_stage = f"{{ duration: '{new_duration}s', target: {target} }}"
        return source[:m.start()] + new_stage + source[m.end():]


@RuleRegistry.register
class K6007MissingTeardown(BaseRule):
    rule_id = "K6007"
    name = "MissingTeardown"
    description = (
        "setup() is exported but teardown() is missing. "
        "Data or state created in setup() (users, orders, sessions) will leak "
        "into the target system. teardown() is the contract for cleanup."
    )
    severity = Severity.WARNING
    frameworks = [Framework.K6]
    tags = ("lifecycle", "correctness")
    fixable = True

    def check(self, ir: ScriptIR) -> list[Violation]:
        if ir.parsed_data.get("has_setup") and not ir.parsed_data.get("has_teardown"):
            return [Violation(
                rule_id=self.rule_id,
                severity=self.severity,
                message="setup() is defined but teardown() is missing. Clean up test data to avoid polluting the target system.",
                location=Location(line=1),
                suggestion="Export a teardown(data) function to clean up any state created in setup().",
                fix_example="export function teardown(data) {\n  // Clean up: delete created users, sessions, orders, etc.\n  http.del(`${BASE_URL}/api/users/${data.userId}`);\n}",
            )]
        return []

    def apply_fix(self, ir: ScriptIR) -> str | None:
        if ir.parsed_data.get("has_teardown"):
            return None
        source = ir.raw_content
        teardown = (
            "\nexport function teardown(data) {\n"
            "  // TODO: clean up resources created in setup()\n"
            "  // e.g. http.del(`${BASE_URL}/api/users/${data.userId}`);\n"
            "}\n"
        )
        return source.rstrip() + "\n" + teardown


@RuleRegistry.register
class K6008HardcodedAuthToken(BaseRule):
    rule_id = "K6008"
    name = "HardcodedAuthToken"
    description = (
        "A hardcoded Authorization header value (Bearer token or Basic credentials) "
        "was detected. Hardcoded tokens expire mid-test, causing all VUs to silently "
        "receive 401 responses."
    )
    severity = Severity.ERROR
    frameworks = [Framework.K6]
    tags = ("security", "parameterization", "correctness")
    tier = "pro"

    def check(self, ir: ScriptIR) -> list[Violation]:
        if ir.parsed_data.get("has_hardcoded_auth_token"):
            return [Violation(
                rule_id=self.rule_id,
                severity=self.severity,
                message="Hardcoded Authorization token detected. Tokens expire — use __ENV.TOKEN or fetch a token in setup().",
                location=Location(line=None),
                suggestion="Load the token from an environment variable (__ENV.TOKEN) or fetch it dynamically in setup().",
                fix_example="const TOKEN = __ENV.TOKEN;\n\nconst res = http.get(url, {\n  headers: { Authorization: `Bearer ${TOKEN}` },\n});",
            )]
        return []


@RuleRegistry.register
class K6009MissingGroup(BaseRule):
    rule_id = "K6009"
    name = "MissingGroup"
    description = (
        "Multiple HTTP calls found but no group() calls. "
        "Without group(), the K6 end-of-test summary reports each URL individually "
        "with no way to see composite response times for logical user journeys."
    )
    severity = Severity.WARNING
    frameworks = [Framework.K6]
    tags = ("structure", "reporting")
    tier = "pro"

    def check(self, ir: ScriptIR) -> list[Violation]:
        http_count = ir.parsed_data.get("http_call_count", 0)
        if http_count > 3 and not ir.parsed_data.get("has_group"):
            return [Violation(
                rule_id=self.rule_id,
                severity=self.severity,
                message=f"Found {http_count} HTTP calls but no group() calls. Group related requests to get composite transaction metrics.",
                location=Location(line=1),
                suggestion="Wrap related requests in group() to report end-to-end transaction times.",
                fix_example="import { group } from 'k6';\n\ngroup('Login Journey', () => {\n  const loginRes = http.post('/login', credentials);\n  check(loginRes, { 'login ok': (r) => r.status === 200 });\n  sleep(1);\n});",
            )]
        return []


@RuleRegistry.register
class K6010MissingRequestTag(BaseRule):
    rule_id = "K6010"
    name = "MissingRequestTag"
    description = (
        "HTTP calls found with no custom tags. Without tags, every URL variation "
        "(UUID in path, query params) becomes its own metric label, causing "
        "cardinality explosion in InfluxDB/Prometheus."
    )
    severity = Severity.WARNING
    frameworks = [Framework.K6]
    tags = ("observability", "configuration")
    fixable = True
    tier = "pro"

    def check(self, ir: ScriptIR) -> list[Violation]:
        http_count = ir.parsed_data.get("http_call_count", 0)
        if http_count > 0 and not ir.parsed_data.get("has_request_tags"):
            return [Violation(
                rule_id=self.rule_id,
                severity=self.severity,
                message="HTTP calls found with no custom tags. Add tags to control metric cardinality and improve result filtering.",
                location=Location(line=None),
                suggestion="Add tags to your http calls to group requests by name and prevent label explosion.",
                fix_example="http.get(url, {\n  tags: { name: 'homepage' },\n});",
            )]
        return []

    def apply_fix(self, ir: ScriptIR) -> str | None:
        source = ir.raw_content
        # Match http.method('url') with no second argument -- single-argument calls
        single_arg_re = re.compile(
            r"""(http\.(?:get|post|put|del|patch|head|options)\s*\()([^,)\n]+)(\s*\))"""
        )
        count = 0

        def replacer(m: re.Match) -> str:  # type: ignore[type-arg]
            nonlocal count
            count += 1
            return f"{m.group(1)}{m.group(2)}, {{ tags: {{ name: 'request_{count}' }} }}{m.group(3)}"

        new_source = single_arg_re.sub(replacer, source)
        return new_source if new_source != source else None


@RuleRegistry.register
class K6011SharedArrayNotUsed(BaseRule):
    rule_id = "K6011"
    name = "SharedArrayNotUsed"
    description = (
        "open() is used to read file data but SharedArray is not used. "
        "Data read inside the default function is re-allocated per VU iteration. "
        "SharedArray allocates once and shares memory across all VUs — critical "
        "for large datasets with many VUs."
    )
    severity = Severity.WARNING
    frameworks = [Framework.K6]
    tags = ("performance", "memory")
    tier = "team"

    def check(self, ir: ScriptIR) -> list[Violation]:
        if ir.parsed_data.get("has_open_call") and not ir.parsed_data.get("uses_shared_array"):
            return [Violation(
                rule_id=self.rule_id,
                severity=self.severity,
                message="open() used without SharedArray. Use SharedArray to share data across VUs and avoid per-iteration allocation.",
                location=Location(line=None),
                suggestion="Wrap file reads in SharedArray to allocate data once and share it across all virtual users.",
                fix_example="import { SharedArray } from 'k6/data';\n\nconst data = new SharedArray('users', function () {\n  return JSON.parse(open('./users.json'));\n});",
            )]
        return []


@RuleRegistry.register
class K6012MissingGracefulStop(BaseRule):
    rule_id = "K6012"
    name = "MissingGracefulStop"
    description = (
        "stages or scenarios defined but no gracefulStop configured. "
        "Without graceful stop, K6 kills VUs mid-request at test end, "
        "producing artificial errors that inflate the final error rate."
    )
    severity = Severity.WARNING
    frameworks = [Framework.K6]
    tags = ("lifecycle", "correctness")
    fixable = True

    def check(self, ir: ScriptIR) -> list[Violation]:
        has_stages = bool(ir.parsed_data.get("stages"))
        has_scenarios = ir.parsed_data.get("has_scenarios", False)
        if (has_stages or has_scenarios) and not ir.parsed_data.get("has_graceful_stop"):
            return [Violation(
                rule_id=self.rule_id,
                severity=self.severity,
                message="Test uses stages/scenarios but has no gracefulStop. VUs will be killed mid-request, inflating error rates.",
                location=Location(line=None),
                suggestion="Add gracefulStop and gracefulRampDown to your options to allow in-flight requests to complete.",
                fix_example="export const options = {\n  stages: [...],\n  gracefulStop: '30s',\n  gracefulRampDown: '30s',\n};",
            )]
        return []

    def apply_fix(self, ir: ScriptIR) -> str | None:
        source = ir.raw_content
        if "gracefulStop" in source or "gracefulRampDown" in source:
            return None
        options_re = re.compile(r"export\s+const\s+options\s*=\s*\{")
        m = options_re.search(source)
        if not m:
            return None
        insert_pos = m.end()
        graceful_block = "\n  gracefulStop: '30s',\n  gracefulRampDown: '30s',"
        return source[:insert_pos] + graceful_block + source[insert_pos:]


@RuleRegistry.register
class K6013ClosedModelOnly(BaseRule):
    rule_id = "K6013"
    name = "ClosedModelOnly"
    description = (
        "Test uses VU-based stages (closed model) with no arrival-rate executor. "
        "Under slow responses, throughput drops because VUs queue. "
        "Consider constant-arrival-rate for throughput-based SLOs."
    )
    severity = Severity.INFO
    frameworks = [Framework.K6]
    tags = ("load-model", "realism")

    def check(self, ir: ScriptIR) -> list[Violation]:
        has_stages = bool(ir.parsed_data.get("stages"))
        has_arrival = ir.parsed_data.get("has_arrival_rate", False)
        if has_stages and not has_arrival:
            return [Violation(
                rule_id=self.rule_id,
                severity=self.severity,
                message="Test uses VU-based (closed model) stages. For RPS-based SLOs, consider constant-arrival-rate executor.",
                location=Location(line=None),
                suggestion="Use a constant-arrival-rate scenario for throughput-based load goals.",
                fix_example="export const options = {\n  scenarios: {\n    api_load: {\n      executor: 'constant-arrival-rate',\n      rate: 100,      // 100 RPS\n      timeUnit: '1s',\n      duration: '5m',\n      preAllocatedVUs: 50,\n    },\n  },\n};",
            )]
        return []


@RuleRegistry.register
class K6014MissingConnectionTimeout(BaseRule):
    rule_id = "K6014"
    name = "MissingConnectionTimeout"
    description = (
        "HTTP calls found with no timeout parameter. "
        "K6 default timeout is 60 seconds per request. Under load with a slow "
        "target, a 60s hang per VU quickly exhausts all VUs."
    )
    severity = Severity.INFO
    frameworks = [Framework.K6]
    tags = ("resilience", "configuration")
    tier = "team"

    def check(self, ir: ScriptIR) -> list[Violation]:
        if ir.parsed_data.get("http_call_count", 0) > 0 and not ir.parsed_data.get("has_http_timeout"):
            return [Violation(
                rule_id=self.rule_id,
                severity=self.severity,
                message="HTTP calls found with no explicit timeout. K6 default is 60s — set an explicit timeout to prevent VU exhaustion.",
                location=Location(line=None),
                suggestion="Set a timeout in http.params or via options.httpDebug.",
                fix_example="const res = http.get(url, {\n  timeout: '10s',  // Fail fast on slow responses\n});",
            )]
        return []


@RuleRegistry.register
class K6015MissingCustomMetrics(BaseRule):
    rule_id = "K6015"
    name = "MissingCustomMetrics"
    description = (
        "Long test script with many HTTP calls but no custom metrics (Trend/Rate/Counter/Gauge). "
        "Tests without custom metrics only capture HTTP stats, missing business outcomes "
        "like checkout success rate or search latency."
    )
    severity = Severity.INFO
    frameworks = [Framework.K6]
    tags = ("observability", "metrics")
    tier = "team"

    def check(self, ir: ScriptIR) -> list[Violation]:
        if ir.parsed_data.get("http_call_count", 0) > 5 and not ir.parsed_data.get("has_custom_metrics"):
            return [Violation(
                rule_id=self.rule_id,
                severity=self.severity,
                message="Script has many HTTP calls but no custom metrics. Add Trend/Rate/Counter to track business KPIs.",
                location=Location(line=None),
                suggestion="Use K6 custom metrics to track business outcomes beyond HTTP response codes.",
                fix_example="import { Trend, Rate } from 'k6/metrics';\n\nconst checkoutDuration = new Trend('checkout_duration');\nconst checkoutSuccess = new Rate('checkout_success');\n\n// In your test:\ncheckoutDuration.add(res.timings.duration);\ncheckoutSuccess.add(res.status === 200);",
            )]
        return []
