"""Prepare complete MCP changes before touching the running configuration."""
import base64
import copy
from application import services as svc, firewall as fw
from application.snapshots import service_snapshot, write_service_data, restore, configuration
from modules.services.models import Service, Address, Filter, Regex
from modules.services.transports import ChainLink, build_class
from modules.services.firewall import code_path
from modules.firewall.models import RuleModel
from mcp_api.auth import MCPError

DIRECTIONS = {"client_to_server": "C", "server_to_client": "S", "both": "B"}


def resolve_service(reference):
    rows = svc.db.query("SELECT service_id FROM services WHERE service_id=? OR name=?", reference, reference)
    if len(rows) != 1:
        raise MCPError("NOT_FOUND" if not rows else "AMBIGUOUS_SERVICE", "Use a unique service name or ID")
    return rows[0]["service_id"]


def resolve_filter(sid, reference):
    rows = svc.db.query("SELECT filter_id FROM filters WHERE service_id=? AND (filter_id=? OR name=?)",
                        sid, reference, reference)
    if len(rows) != 1:
        raise MCPError("NOT_FOUND" if not rows else "AMBIGUOUS_FILTER", "Use a unique filter name or ID",
                       candidates=[r["filter_id"] for r in rows])
    return rows[0]["filter_id"]


def check_chain(value):
    srv = Service.from_dict(value["service"])
    srv.addresses = [Address.from_dict(r) for r in value["addresses"]]
    links = [ChainLink(Filter.from_dict(row),
                       regexes=[Regex.from_dict(r) for r in value["regexes"] if r["filter_id"] == row["filter_id"]],
                       code_path=code_path(row["filter_id"]),
                       functions=[r for r in value["functions"] if r["filter_id"] == row["filter_id"]])
             for row in value["filters"]]
    build_class(srv.transport).check(srv, links)


async def prepare_service(before, spec):
    after = copy.deepcopy(before)
    sid = before["service_id"]
    changes = spec.model_dump(exclude_none=True, exclude={"addresses"})
    if before["service"] is None:
        if not spec.name or not spec.addresses:
            raise MCPError("INVALID_ARGUMENT", "Creating a service requires name and addresses")
        row = {"service_id": sid, "name": spec.name, "status": "stop", "proto": "tcp", "transport": "proxy",
               "fail_open": 1, "max_connections": 0, "over_limit_forwards": 0, "first_byte_timeout": 0,
               "tls_cert": None, "tls_key": None}
    else:
        row = dict(before["service"])
    row.update(changes)
    row["name"] = svc.refactor_name(row["name"])
    if svc.db.query("SELECT service_id FROM services WHERE name=? AND service_id<>?", row["name"], sid):
        raise MCPError("NAME_CONFLICT", "A service with that name already exists")
    svc._check_tls_material(row["tls_cert"], row["tls_key"])
    if row["proto"] in {"tls", "quic", "http"} and not (row["tls_cert"] and row["tls_key"]):
        raise MCPError("INVALID_ARGUMENT", "This protocol needs a certificate and private key")
    if spec.addresses is not None:
        old = {r["address_id"]: r for r in before["addresses"]}
        addresses = []
        for item in spec.addresses:
            form = svc.AddressForm(**item.model_dump(exclude={"address_id"}))
            ip = svc.parse_ip_or_int(form.ip_int)
            edge = svc._address_edge(row["proto"], form)
            aid = item.address_id
            if aid and aid not in old:
                raise MCPError("NOT_FOUND", "address_id does not belong to this service")
            if not aid:
                aid = next((a["address_id"] for a in before["addresses"]
                            if (a["ip_int"], a["port"], a["edge"]) == (ip, form.port, edge)), svc.gen_id())
            addresses.append({"address_id": aid, "service_id": sid, "ip_int": ip, "port": form.port,
                              "proto": svc.L4.l4_of(edge), "edge": edge,
                              "target_port": svc._address_target(row["transport"], form),
                              "upstream": svc._address_upstream(row["proto"], edge, form),
                              "proxy_ip": svc._hijack_ip(row["transport"], ip, form),
                              "proxy_port": svc._hijack_port(row["transport"], form)})
    else:
        addresses = [svc._address_under(a, row["proto"]) for a in before["addresses"]]
    seen, proxy_seen, ids = set(), set(), set()
    for address in addresses:
        key = (address["ip_int"], address["port"], address["proto"])
        if key in seen or address["address_id"] in ids or svc.db.query(
            "SELECT address_id FROM service_addresses WHERE ip_int=? AND port=? AND proto=? AND service_id<>?", *key, sid
        ):
            raise MCPError("ADDRESS_CONFLICT", "An address is duplicated or already belongs to another service")
        seen.add(key)
        ids.add(address["address_id"])
        if row["transport"] != "proxy" and address["target_port"]:
            raise MCPError("UNSUPPORTED_COMBINATION", "Published target ports require the proxy transport")
        if row["transport"] == "external":
            if not svc.is_ip_parse(address["ip_int"]) or not address["proxy_port"]:
                raise MCPError("UNSUPPORTED_COMBINATION", "External transport requires a concrete IP and proxy port")
            key = (address["proxy_ip"], address["proxy_port"])
            if key in proxy_seen or svc.db.query(
                "SELECT address_id FROM service_addresses WHERE proxy_ip=? AND proxy_port=? AND service_id<>?", *key, sid
            ):
                raise MCPError("ADDRESS_CONFLICT", "This external proxy endpoint is already in use")
            proxy_seen.add(key)
        else:
            address.update(proxy_ip=None, proxy_port=None)
    after.update(service=row, addresses=sorted(addresses, key=lambda a: a["address_id"]))
    check_chain(after)
    if row["proto"] == "udp":
        for code in after["code"].values():
            checked = await svc.validate_code(code)
            if not checked["ok"] or svc.needs_a_stream(checked.get("models", [])):
                raise MCPError("UNSUPPORTED_COMBINATION", "A Python filter cannot load on this UDP service", check=checked)
    return after


async def commit_service(before, after):
    sid = after["service_id"]
    if configuration(before) == configuration(after):
        return {"service_id": sid, "changed": False}
    old, new = before["service"], after["service"]
    old_addresses = {a["address_id"]: a for a in before["addresses"]}
    new_addresses = {a["address_id"]: a for a in after["addresses"]}
    append_only = old and all(k == "name" or old.get(k) == v for k, v in new.items()) and all(
        new_addresses.get(k) == v for k, v in old_addresses.items())
    if append_only:
        write_service_data(after)
        manager = svc.firewall.get(sid)
        manager.srv.name = new["name"]
        for aid in new_addresses.keys() - old_addresses.keys():
            await manager.address_added(aid, lambda: write_service_data(before))
        await svc.refresh_frontend()
    else:
        await restore(after)
    return {"service_id": sid, "observed_state": "running" if svc.firewall.get(sid).active else "stopped"}


async def prepare_filter(before, filter_id, spec):
    after = copy.deepcopy(before)
    sid = before["service_id"]
    old = next((f for f in after["filters"] if f["filter_id"] == filter_id), None)
    if old is None:
        if spec.kind is None:
            raise MCPError("INVALID_ARGUMENT", "Creating a filter requires kind")
        row = {"filter_id": filter_id, "service_id": sid, "position": len(after["filters"]),
               "kind": spec.kind, "name": spec.name or spec.kind, "active": 1, "blocked": 0, "proto": "tcp"}
        after["filters"].append(row)
    else:
        row = old
        if spec.kind is not None and spec.kind != row["kind"]:
            raise MCPError("INVALID_ARGUMENT", "An existing filter cannot change kind")
    if spec.name is not None:
        row["name"] = svc.refactor_name(spec.name)
    if spec.active is not None:
        row["active"] = int(spec.active)
    if row["kind"] == "regex":
        if spec.code is not None or spec.functions is not None:
            raise MCPError("INVALID_ARGUMENT", "Code/functions belong to a pyfilter")
        if spec.patterns is not None:
            prior = {r["regex_id"]: r for r in after["regexes"] if r["filter_id"] == filter_id}
            patterns, seen = [], set()
            for item in spec.patterns:
                ok, why = await svc.check_pattern(item.pattern, item.case_sensitive)
                if not ok:
                    raise MCPError("FILTER_COMPILE_ERROR", why, pattern=item.pattern)
                mode = DIRECTIONS[item.direction]
                encoded = base64.b64encode(item.pattern.encode()).decode()
                key = (encoded, mode, item.case_sensitive)
                if key in seen:
                    raise MCPError("INVALID_ARGUMENT", "Duplicate pattern in this filter")
                seen.add(key)
                rid = item.regex_id
                if rid and rid not in prior:
                    raise MCPError("NOT_FOUND", "regex_id does not belong to this filter")
                if not rid:
                    rid = next((r["regex_id"] for r in prior.values() if
                                (r["regex"], r["mode"], bool(r["case_sensitive"])) == key), svc.gen_id())
                if any(r["regex_id"] == rid for r in patterns):
                    raise MCPError("INVALID_ARGUMENT", "Duplicate regex_id")
                patterns.append({"regex_id": rid, "filter_id": filter_id, "regex": encoded, "mode": mode,
                                 "case_sensitive": int(item.case_sensitive), "active": int(item.active),
                                 "blocked": prior.get(rid, {}).get("blocked", 0)})
            after["regexes"] = [r for r in after["regexes"] if r["filter_id"] != filter_id] + patterns
    else:
        if spec.patterns is not None:
            raise MCPError("INVALID_ARGUMENT", "Patterns belong to a regex filter")
        source = spec.code if spec.code is not None else after["code"].get(filter_id, "")
        checked = await svc.validate_code(source)
        if not checked["ok"]:
            raise MCPError("FILTER_COMPILE_ERROR", "Python filter could not load", check=checked)
        if after["service"]["proto"] == "udp" and svc.needs_a_stream(checked.get("models", [])):
            raise MCPError("UNSUPPORTED_COMBINATION", "This filter requires a stream; the service carries UDP datagrams")
        row["proto"] = checked.get("proto") or "tcp"
        after["code"][filter_id] = source
        prior = {r["name"]: r for r in after["functions"] if r["filter_id"] == filter_id}
        names = checked.get("filters", [])
        if spec.functions and set(spec.functions) - set(names):
            raise MCPError("NOT_FOUND", "Unknown Python function names", functions=sorted(set(spec.functions) - set(names)))
        after["functions"] = [r for r in after["functions"] if r["filter_id"] != filter_id] + [
            {"filter_id": filter_id, "name": name, "position": i,
             "active": int((spec.functions or {}).get(name, prior.get(name, {}).get("active", True))),
             "blocked": prior.get(name, {}).get("blocked", 0)} for i, name in enumerate(names)]
    if spec.position is not None:
        if spec.position >= len(after["filters"]):
            raise MCPError("INVALID_ARGUMENT", "position must be within the resulting chain")
        after["filters"].remove(row)
        after["filters"].insert(spec.position, row)
    for i, item in enumerate(after["filters"]):
        item["position"] = i
    after["regexes"].sort(key=lambda r: (r["filter_id"], r["regex_id"]))
    check_chain(after)
    return after


async def commit_chain(before, after):
    if configuration(before) == configuration(after):
        return {"service_id": after["service_id"], "changed": False}
    write_service_data(after)
    await svc._apply_chain(after["service_id"], lambda: write_service_data(before))
    # A changed matcher's old history belongs to its filter, as in the REST editor.
    previous = {r["regex_id"]: r for r in before["regexes"]}
    for row in after["regexes"]:
        old = previous.get(row["regex_id"])
        if old and old["regex"] != row["regex"]:
            svc.db.query("UPDATE regexes SET blocked=0 WHERE regex_id=?", row["regex_id"])
            svc.stats.disown_rule(svc.db, row["filter_id"], row["regex_id"])
    removed = {f["filter_id"] for f in before["filters"]} - {f["filter_id"] for f in after["filters"]}
    for fid in removed:
        svc.stats.forget_filter(svc.db, fid)
    await svc.refresh_frontend()
    return {"service_id": after["service_id"], "chain": [r["filter_id"] for r in after["filters"]]}


def prepare_firewall(before, patch):
    after = copy.deepcopy(before)
    if patch.rules is not None:
        rows = [fw.parse_and_check_rule(RuleModel(**r.model_dump(exclude={"id"}))).model_dump(mode="json") for r in patch.rules]
        seen = set()
        for row in rows:
            key = tuple(row[k] for k in ("proto", "src", "dst", "port_src_from", "port_src_to",
                                         "port_dst_from", "port_dst_to", "mode", "table"))
            if key in seen:
                raise MCPError("INVALID_ARGUMENT", "Duplicate firewall rule")
            seen.add(key)
        after["rules"] = rows
    if patch.policy is not None:
        after["policy"] = patch.policy
    if patch.settings:
        after["settings"].update(patch.settings)
    return after