"""Closed tool schemas and deterministic plans for three supported questions."""
from typing import Annotated, Literal
import json
from pydantic import Field, StrictInt, ValidationError
from service_analysis.analysis import Period, Grouping
from service_analysis.schemas import Contract

MAX_STEPS = 6
MAX_PLAN_BYTES = 32768


class SummaryCall(Contract):
    operation: Literal["summarize"]
    period: Period
    grouping: Grouping = "overall"


class ComparisonCall(Contract):
    operation: Literal["compare"]
    baseline: Period
    comparison: Period
    grouping: Grouping = "overall"


class DecompositionCall(Contract):
    operation: Literal["decompose"]
    baseline: Period
    comparison: Period


class WorkloadCall(Contract):
    operation: Literal["workload"]
    opening_period: Period


ToolCall = Annotated[SummaryCall | ComparisonCall | DecompositionCall | WorkloadCall,
                     Field(discriminator="operation")]


class ProposedStep(Contract):
    call: ToolCall
    depends_on: tuple[StrictInt, ...] = Field(default=(), max_length=MAX_STEPS)


class AnalysisPlan(Contract):
    steps: tuple[ProposedStep, ...] = Field(min_length=1, max_length=MAX_STEPS)


class CompiledStep(Contract):
    step_id: str
    call: ToolCall
    depends_on: tuple[str, ...]


class PlanRejected(ValueError):
    def __init__(self, code):
        self.code = code
        super().__init__(code)


def _pairs(values):
    result = {}
    for key, value in values:
        if key in result:
            raise ValueError("Duplicate JSON key")
        result[key] = value
    return result


def _nonfinite(value):
    raise ValueError("Nonfinite JSON value")


def parse_plan(raw):
    try:
        if isinstance(raw, AnalysisPlan):
            raw = raw.model_dump(mode="json")
        if isinstance(raw, str):
            if len(raw.encode()) > MAX_PLAN_BYTES:
                raise PlanRejected("PLAN_TOO_LARGE")
            raw = json.loads(raw, object_pairs_hook=_pairs, parse_constant=_nonfinite)
        if not isinstance(raw, dict):
            raise ValueError("Expected plan object")
        if len(json.dumps(raw, allow_nan=False).encode()) > MAX_PLAN_BYTES:
            raise PlanRejected("PLAN_TOO_LARGE")
        return AnalysisPlan.model_validate(raw)
    except PlanRejected:
        raise
    except (ValueError, TypeError, RecursionError, ValidationError):
        raise PlanRejected("INVALID_PLAN_SCHEMA") from None


def call_key(call):
    return json.dumps(call.model_dump(mode="json"), sort_keys=True, separators=(",", ":"))


def compile_plan(raw, coverage, prior=()):
    """Validate the entire round before assigning executable steps.

    Dependencies use zero-based positions in the entire run, including prior
    rounds. The application translates positions into stable step IDs.
    """
    plan = parse_plan(raw)
    offset = len(prior)
    if offset + len(plan.steps) > MAX_STEPS:
        raise PlanRejected("STEP_BUDGET_EXCEEDED")
    seen = {call_key(step.call) for step in prior}
    compiled = []
    for index, proposed in enumerate(plan.steps, offset):
        dependencies = proposed.depends_on
        if len(set(dependencies)) != len(dependencies) or any(d < 0 or d >= index for d in dependencies):
            raise PlanRejected("INVALID_DEPENDENCY")
        call = proposed.call
        key = call_key(call)
        if key in seen:
            raise PlanRejected("DUPLICATE_CALL")
        seen.add(key)
        if isinstance(call, WorkloadCall):
            periods = (call.opening_period,)
            lower, upper = coverage.operations_start, coverage.operations_end
        else:
            periods = (call.period,) if isinstance(call, SummaryCall) else (call.baseline, call.comparison)
            lower, upper = coverage.resolution_start, coverage.resolution_end
        if any(p.start < lower or p.end > upper for p in periods):
            raise PlanRejected("PERIOD_OUTSIDE_COVERAGE")
        if isinstance(call, (ComparisonCall, DecompositionCall)) and call.baseline.end >= call.comparison.start:
            raise PlanRejected("INVALID_COMPARISON_PERIODS")
        compiled.append(CompiledStep(step_id=f"step-{index + 1:03d}", call=call,
                        depends_on=tuple(f"step-{d + 1:03d}" for d in dependencies)))
    return tuple(compiled)


def deterministic_plan(question, baseline: Period, comparison: Period,
                       opening_period: Period | None = None) -> AnalysisPlan:
    compare = ComparisonCall(operation="compare", baseline=baseline, comparison=comparison)
    steps = [ProposedStep(call=compare)]
    if question == "resolution_change":
        pass
    elif question == "segments":
        steps.extend(ProposedStep(call=ComparisonCall(operation="compare", baseline=baseline,
            comparison=comparison, grouping=group), depends_on=(0,)) for group in ("category", "region"))
    elif question == "drivers":
        if opening_period is None:
            raise ValueError("Drivers question requires an explicit opening period")
        steps.extend([
            ProposedStep(call=DecompositionCall(operation="decompose", baseline=baseline,
                         comparison=comparison), depends_on=(0,)),
            ProposedStep(call=WorkloadCall(operation="workload", opening_period=opening_period)),
        ])
    else:
        raise ValueError("Unsupported question")
    return AnalysisPlan(steps=tuple(steps))