from __future__ import annotations

import json
import shlex
from pathlib import Path

from .models import Finding, PolicyRule

ANY_VALUES = {"all", "any", "0.0.0.0/0"}
UTM_KEYS = {"ips-sensor", "av-profile", "webfilter-profile", "application-list", "waf-profile", "dnsfilter-profile"}


def _parse_set_values(line: str) -> tuple[str, list[str]] | None:
    try:
        parts = shlex.split(line)
    except ValueError:
        parts = line.split()
    if len(parts) < 3 or parts[0] != "set":
        return None
    return parts[1], parts[2:]


def parse_policy_config(text: str) -> list[PolicyRule]:
    in_policy = False
    current_id: str | None = None
    current_settings: dict[str, list[str]] = {}
    policies: list[PolicyRule] = []

    for raw_line in text.splitlines():
        line = raw_line.strip()
        if line == "config firewall policy":
            in_policy = True
            continue
        if in_policy and line == "end":
            if current_id is not None:
                policies.append(PolicyRule(current_id, current_settings))
            break
        if not in_policy:
            continue
        if line.startswith("edit "):
            if current_id is not None:
                policies.append(PolicyRule(current_id, current_settings))
            current_id = line.split(maxsplit=1)[1].strip('"')
            current_settings = {}
            continue
        if line == "next":
            if current_id is not None:
                policies.append(PolicyRule(current_id, current_settings))
            current_id = None
            current_settings = {}
            continue
        parsed = _parse_set_values(line)
        if current_id is not None and parsed:
            key, values = parsed
            current_settings[key] = values

    return policies


def _normalize_policy_value(value: object) -> list[str]:
    if value is None:
        return []
    if isinstance(value, list):
        normalized: list[str] = []
        for item in value:
            if isinstance(item, dict):
                if "name" in item:
                    normalized.append(str(item["name"]))
                elif "q_origin_key" in item:
                    normalized.append(str(item["q_origin_key"]))
                else:
                    normalized.append(json.dumps(item, sort_keys=True))
            else:
                normalized.append(str(item))
        return normalized
    if isinstance(value, dict):
        if "name" in value:
            return [str(value["name"])]
        if "q_origin_key" in value:
            return [str(value["q_origin_key"])]
        return [json.dumps(value, sort_keys=True)]
    return [str(value)]


def parse_policy_json(text: str) -> list[PolicyRule]:
    data = json.loads(text)
    raw_policies = data.get("results", data) if isinstance(data, dict) else data
    if not isinstance(raw_policies, list):
        raise ValueError("policy JSON must be a list or contain a results list")

    policies: list[PolicyRule] = []
    for index, raw_policy in enumerate(raw_policies, start=1):
        if not isinstance(raw_policy, dict):
            continue
        policy_id = str(raw_policy.get("policyid", raw_policy.get("policy_id", raw_policy.get("id", index))))
        settings = {str(key): _normalize_policy_value(value) for key, value in raw_policy.items()}
        policies.append(PolicyRule(policy_id, settings))
    return policies


def read_policies(path: str | Path) -> list[PolicyRule]:
    policy_path = Path(path)
    if not policy_path.exists():
        raise FileNotFoundError(
            f"{policy_path} does not exist. Fetch policies first with: "
            f"fgai fetch-policies --output exports/policies.json"
        )

    text = policy_path.read_text(encoding="utf-8", errors="replace")
    stripped = text.lstrip()
    if stripped.startswith("{") or stripped.startswith("["):
        return parse_policy_json(text)
    return parse_policy_config(text)


def audit_policies(policies: list[PolicyRule]) -> list[Finding]:
    findings: list[Finding] = []
    for policy in policies:
        if policy.first("status", "enable") == "disable":
            continue

        action = policy.first("action")
        srcaddr = {value.lower() for value in policy.values("srcaddr")}
        dstaddr = {value.lower() for value in policy.values("dstaddr")}
        service = {value.lower() for value in policy.values("service")}
        policy_ref = f"policy {policy.policy_id}"

        if action == "accept" and srcaddr & ANY_VALUES and dstaddr & ANY_VALUES and service & ANY_VALUES:
            findings.append(
                Finding(
                    severity="critical",
                    title="Broad allow policy",
                    detail="Accepts traffic from any source to any destination for any service.",
                    reference=policy_ref,
                )
            )

        if action == "accept" and policy.first("logtraffic", "disable") in {"disable", "utm"}:
            findings.append(
                Finding(
                    severity="medium",
                    title="Insufficient traffic logging",
                    detail="Accepted traffic is not fully logged, reducing investigation value.",
                    reference=policy_ref,
                )
            )

        has_utm = policy.first("utm-status") == "enable" or any(policy.values(key) for key in UTM_KEYS)
        if action == "accept" and not has_utm:
            findings.append(
                Finding(
                    severity="high",
                    title="Accepted traffic lacks UTM inspection",
                    detail="No IPS, AV, web filter, application, DNS, or WAF profile is attached.",
                    reference=policy_ref,
                )
            )

    return findings
