from __future__ import annotations

import json
from datetime import datetime
from typing import Dict, Any

from sqlmodel import SQLModel, Field
from sqlalchemy import Column, Text


class PolicyRow(SQLModel, table=True):
    __tablename__ = "policy_store"
    __table_args__ = {"extend_existing": True}

    account_id: str = Field(primary_key=True)
    segment: str = Field(primary_key=True)
    policy_json: str = Field(sa_column=Column(Text), default="{}")
    version: int = Field(default=1)
    updated_at: datetime = Field(default_factory=datetime.utcnow)


class PolicyStore:
    """Intelligence-owned policy store.

    Services must only READ via IntelligenceClient.
    Learning pipeline is the only writer via set_policy.
    """

    def __init__(self, db):
        self.db = db

    async def get_policy(self, account_id: str, segment: str) -> Dict[str, Any]:
        row = await self.db.get(PolicyRow, (str(account_id), str(segment)))
        if not row:
            # Reasonable defaults owned by Intelligence
            return {
                "suggest": 0.70,
                "min_gap": 0.10,
                "auto_link": 0.80,
                "weights": {"fuzzy_a": 1.0, "fuzzy_b": 1.0, "exact": 1.0},
                "updated_at": None,
                "version": 1,
            }
        try:
            data = json.loads(row.policy_json or "{}")
        except Exception:
            data = {}
        data.setdefault("version", row.version)
        data.setdefault(
            "updated_at", row.updated_at.isoformat() if row.updated_at else None
        )
        return data

    async def set_policy(
        self, account_id: str, segment: str, policy: Dict[str, Any]
    ) -> Dict[str, Any]:
        row = await self.db.get(PolicyRow, (str(account_id), str(segment)))
        if not row:
            row = PolicyRow(
                account_id=str(account_id),
                segment=str(segment),
                policy_json=json.dumps(policy or {}),
                version=int(policy.get("version") or 1),
            )
            self.db.add(row)
        else:
            row.policy_json = json.dumps(policy or {})
            row.version = (
                int(policy.get("version") or (row.version or 1)) + 0
            )  # keep as provided
            row.updated_at = datetime.utcnow()
        await self.db.commit()
        return await self.get_policy(account_id, segment)
