#!/usr/bin/env python3
"""Update context_policy.py with strategy configuration support."""

import sys
from pathlib import Path

# Read the file
policy_file = Path("src/cleveragents/domain/models/core/context_policy.py")
content = policy_file.read_text()

# Update the imports to include Any
content = content.replace(
    "from typing import TYPE_CHECKING", "from typing import TYPE_CHECKING, Any"
)

# Add VALID_STRATEGIES constant after VALID_PHASES
old_phases = (
    "VALID_PHASES: frozenset[str] = frozenset("
    '{"default", "strategize", "execute", "apply"})'
)
new_phases = (
    "VALID_PHASES: frozenset[str] = frozenset("
    '{"default", "strategize", "execute", "apply"})\n\n'
    "VALID_STRATEGIES: frozenset[str] = frozenset({\n"
    '    "basic",\n'
    '    "semantic",\n'
    '    "relevance_scoring",\n'
    '    "adaptive",\n'
    '    "fusion",\n'
    "})"
)
content = content.replace(old_phases, new_phases)

# Update the docstring for ProjectContextPolicy to mention strategy
old_docstring = (
    "class ProjectContextPolicy(BaseModel):\n"
    '    """Controls what context is available during each ACMS phase.\n'
    "\n"
    "    Uses view inheritance: ``default`` → ``strategize`` →\n"
    "    ``execute`` → ``apply``.  Each phase can override or inherit\n"
    "    from its parent.\n"
    "\n"
    "    An empty ``ProjectContextPolicy()`` defaults to including\n"
    "    everything (the ``default_view`` has empty include lists which\n"
    '    means "all").\n'
    '    """'
)

new_docstring = (
    "class ProjectContextPolicy(BaseModel):\n"
    '    """Controls what context is available during each ACMS phase.\n'
    "\n"
    "    Uses view inheritance: ``default`` → ``strategize`` →\n"
    "    ``execute`` → ``apply``.  Each phase can override or inherit\n"
    "    from its parent.\n"
    "\n"
    "    An empty ``ProjectContextPolicy()`` defaults to including\n"
    "    everything (the ``default_view`` has empty include lists which\n"
    '    means "all").\n'
    "\n"
    "    Optionally specifies a context assembly strategy and its\n"
    "    configuration parameters.\n"
    '    """'
)

content = content.replace(old_docstring, new_docstring)

# Add strategy and strategy_config fields before the resolve_view method
old_fields = (
    "    apply_view: ContextView | None = Field(\n"
    "        default=None,\n"
    '        description=("Overrides for Apply (inherits from execute if None)"),\n'
    "    )\n"
    "\n"
    "    def resolve_view(self, phase: str) -> ContextView:"
)

error_msg = "f\"Invalid strategy '{v}': must be one of {sorted(VALID_STRATEGIES)}\""

new_fields = (
    "    apply_view: ContextView | None = Field(\n"
    "        default=None,\n"
    '        description=("Overrides for Apply (inherits from execute if None)"),\n'
    "    )\n"
    "    strategy: str | None = Field(\n"
    "        default=None,\n"
    "        description=(\n"
    '            "Context assembly strategy name. "\n'
    '            "Valid values: basic, semantic, relevance_scoring, adaptive, fusion"\n'
    "        ),\n"
    "    )\n"
    "    strategy_config: dict[str, Any] | None = Field(\n"
    "        default=None,\n"
    '        description="Strategy-specific configuration parameters",\n'
    "    )\n"
    "\n"
    '    @field_validator("strategy")\n'
    "    @classmethod\n"
    "    def _validate_strategy(\n"
    "        cls: type[ProjectContextPolicy],\n"
    "        v: str | None,\n"
    "    ) -> str | None:\n"
    '        """Validate that strategy name is in the list of valid strategies."""\n'
    "        if v is not None and v not in VALID_STRATEGIES:\n"
    "            raise ValueError(\n"
    "                " + error_msg + "\n"
    "            )\n"
    "        return v\n"
    "\n"
    "    def resolve_view(self, phase: str) -> ContextView:"
)

content = content.replace(old_fields, new_fields)

# Write the updated content
policy_file.write_text(content)

print("Updated context_policy.py successfully")
sys.exit(0)