"""
CostPolicy — the main SDK class.

Usage:
    from costwise_mcp import CostPolicy

    policy = CostPolicy(api_key="cw_...")

    decision = policy.optimize(
        prompt="Translate 'hello' to French",
        model="gpt-4",
    )
    # decision.recommended_model = "gpt-4o-mini"
    # decision.max_tokens = 256
    # decision.estimated_cost = 0.000045

    # Call your LLM with the optimized params
    response = openai.chat.completions.create(
        model=decision.recommended_model,
        max_tokens=decision.max_tokens,
        messages=[{"role": "user", "content": prompt}],
    )

    # Report actual usage for analytics (async, non-blocking)
    policy.report(decision, actual_tokens=response.usage.total_tokens)
"""

from typing import Optional
from costwise_mcp.models import Decision, PolicyConfig, TaskComplexity, UsageReport
from costwise_mcp.pricing import (
    PRICING,
    DOWNGRADE_MAP,
    detect_provider,
    estimate_cost,
    get_cheapest_alternative,
)
from costwise_mcp.classifier import classify_task, estimate_tokens
from costwise_mcp.telemetry import TelemetryClient, NoopTelemetry


DEFAULT_BACKEND = "https://app.cost-wise.dev"


class CostPolicy:
    """
    AI cost optimization engine.

    Analyzes prompts, classifies complexity, recommends cheaper models,
    and limits tokens — all locally, no network calls in the hot path.

    Requires a valid API key for optimization. Without a key, optimize()
    returns the original model unchanged (passthrough mode).

    Args:
        api_key: CostWise API key (required for optimization + telemetry).
        backend_url: CostWise backend URL. Default: https://app.cost-wise.dev
        config: PolicyConfig for custom token limits, budgets, etc.
        project_id: Optional project identifier for grouping analytics.
    """

    def __init__(
        self,
        api_key: Optional[str] = None,
        backend_url: str = DEFAULT_BACKEND,
        config: Optional[PolicyConfig] = None,
        project_id: Optional[str] = None,
    ):
        self.config = config or PolicyConfig()
        self.project_id = project_id
        self._api_key = api_key
        self._backend_url = backend_url
        self._activated = False
        self._last_license_check = 0.0
        self._license_check_interval = 3600.0  # re-verify every hour

        # Validate API key format
        if api_key and api_key.startswith("cw_") and len(api_key) > 10:
            self._activated = True
        elif api_key:
            import warnings
            warnings.warn(
                "Invalid CostWise API key format. Expected 'cw_...' "
                "Optimization is disabled — passthrough mode.",
                stacklevel=2,
            )

        # Initialize telemetry (async, non-blocking)
        if self._activated and self.config.telemetry_enabled:
            self._telemetry = TelemetryClient(
                endpoint=backend_url,
                api_key=api_key,
                batch_size=self.config.telemetry_batch_size,
                flush_interval=self.config.telemetry_flush_interval,
            )
        else:
            self._telemetry = NoopTelemetry()

    def _check_license(self) -> bool:
        """
        Periodic license check — verifies the API key is still valid.
        Runs at most once per hour in a background thread.
        NEVER blocks the caller — returns cached state immediately.
        """
        import time
        import threading
        now = time.time()

        if now - self._last_license_check < self._license_check_interval:
            return self._activated

        # Mark as checked NOW to prevent concurrent checks
        self._last_license_check = now

        if not self._api_key:
            return False

        # Run verification in background thread — never blocks optimize()
        def _verify():
            try:
                import urllib.request
                import json
                req = urllib.request.Request(
                    f"{self._backend_url}/api/mcp/verify",
                    headers={"Authorization": f"Bearer {self._api_key}"},
                    method="GET",
                )
                resp = urllib.request.urlopen(req, timeout=5)
                data = json.loads(resp.read())
                self._activated = data.get("valid", False)
            except Exception:
                # Network error — keep last known state
                pass

        threading.Thread(target=_verify, daemon=True).start()
        return self._activated

    def _get_team_policy(self, team_id: str) -> Optional[dict]:
        """
        Fetch team policy from backend. Cached for 5 minutes.
        Returns None if no policy exists or network is unreachable.
        """
        import time
        now = time.time()
        cache_key = f"policy:{team_id}"

        # Check cache
        if not hasattr(self, "_policy_cache"):
            self._policy_cache: dict = {}
        cached = self._policy_cache.get(cache_key)
        if cached and now - cached["_fetched_at"] < 300:  # 5 min TTL
            return cached

        # Fetch from backend (non-blocking — return stale cache on failure)
        try:
            import urllib.request
            import json
            req = urllib.request.Request(
                f"{self._backend_url}/api/mcp/policies?team={team_id}",
                headers={"Authorization": f"Bearer {self._api_key}"},
                method="GET",
            )
            resp = urllib.request.urlopen(req, timeout=3)
            data = json.loads(resp.read())
            data["_fetched_at"] = now
            self._policy_cache[cache_key] = data
            return data
        except Exception:
            # Return stale cache if available
            return cached if cached else None

    def optimize(
        self,
        prompt: str,
        model: str,
        task_type: Optional[TaskComplexity] = None,
        max_budget: Optional[float] = None,
        team: Optional[str] = None,
    ) -> Decision:
        """
        Optimize an LLM request — the core SDK function.

        Runs entirely locally (no network calls). Returns a Decision object
        with the recommended model, max_tokens, and estimated cost.

        Args:
            prompt: The prompt text (used for token estimation + complexity classification).
                    Never sent to the backend — only metadata is reported.
            model: The model the user intends to use (e.g. "gpt-5.4").
            task_type: Override automatic complexity classification.
            max_budget: Max cost in USD for this request. None = no limit.
            team: Team identifier for policy enforcement (e.g. "interns", "ai-team").
                  Maps to MCPTeamPolicy.teamId on the backend.
        
        Runs entirely locally (no network calls). Returns a Decision object
        with the recommended model, max_tokens, and estimated cost.

        Args:
            prompt: The prompt text (used for token estimation + complexity classification).
                    Never sent to the backend — only metadata is reported.
            model: The model the user intends to use (e.g. "gpt-4").
            task_type: Override automatic complexity classification.
            max_budget: Max cost in USD for this request. None = no limit.

        Returns:
            Decision with recommended_model, max_tokens, estimated_cost, etc.
        """
        # ── Step 0: License check ─────────────────────────────────
        if not self._check_license():
            # Passthrough mode — return original model with no optimization
            input_tokens = estimate_tokens(prompt)
            complexity = task_type or classify_task(prompt)
            max_tokens = self.config.get_max_tokens(complexity)
            cost = estimate_cost(model, input_tokens, max_tokens)
            return Decision(
                recommended_model=model,
                max_tokens=max_tokens,
                estimated_cost=cost,
                original_model=model,
                input_tokens=input_tokens,
                complexity=complexity,
                allowed=True,
                savings_pct=0.0,
                message="CostWise optimization disabled — no valid API key. "
                        "Get one at https://app.cost-wise.dev",
                original_estimated_cost=cost,
                estimated_savings=0.0,
            )

        # ── Step 1: Estimate input tokens ─────────────────────────
        input_tokens = estimate_tokens(prompt)

        # ── Step 2: Classify task complexity ──────────────────────
        complexity = task_type or classify_task(prompt)

        # ── Step 3: Determine max output tokens ──────────────────
        max_tokens = self.config.get_max_tokens(complexity)

        # ── Step 3.5: Apply team policy (if team is specified) ────
        team_policy = self._get_team_policy(team) if team else None
        policy_message = ""

        if team_policy:
            # Override max_tokens if policy sets a limit
            if team_policy.get("max_tokens_per_request"):
                max_tokens = min(max_tokens, team_policy["max_tokens_per_request"])

            # Check model against tier restriction
            model_tier = PRICING.get(model, {}).get("tier", 1)
            max_tier = team_policy.get("max_model_tier", 3)

            if model_tier < max_tier:
                # Model is above the team's allowed tier — force downgrade
                downgraded = DOWNGRADE_MAP.get(model)
                while downgraded and PRICING.get(downgraded, {}).get("tier", 1) < max_tier:
                    next_down = DOWNGRADE_MAP.get(downgraded)
                    if not next_down:
                        break
                    downgraded = next_down
                if downgraded and PRICING.get(downgraded, {}).get("tier", 1) >= max_tier:
                    policy_message = f"Team '{team}' restricted to tier {max_tier} models. {model} → {downgraded}"
                    model = downgraded

            # Check blocked models
            blocked = team_policy.get("blocked_models", [])
            if model in blocked:
                alt = get_cheapest_alternative(model, detect_provider(model))
                if alt not in blocked:
                    policy_message = f"Model '{model}' blocked for team '{team}'. Using {alt}"
                    model = alt

            # Check allowed models (if set, overrides tier)
            allowed = team_policy.get("allowed_models", [])
            if allowed and model not in allowed:
                # Pick the first allowed model from the same provider
                provider = detect_provider(model)
                for am in allowed:
                    if detect_provider(am) == provider:
                        policy_message = f"Team '{team}' restricted to allowed models. {model} → {am}"
                        model = am
                        break

            # Check budget
            budget_limit = team_policy.get("monthly_budget_usd")
            current_spend = team_policy.get("current_month_spend", 0)
            if budget_limit and max_budget is None:
                remaining = budget_limit - current_spend
                if remaining > 0:
                    max_budget = remaining
                else:
                    # Budget exceeded — block
                    return Decision(
                        recommended_model=model,
                        max_tokens=max_tokens,
                        estimated_cost=0,
                        original_model=model,
                        input_tokens=input_tokens,
                        complexity=complexity,
                        allowed=False,
                        savings_pct=0,
                        message=f"Team '{team}' monthly budget exceeded (${current_spend:.2f} / ${budget_limit:.2f})",
                        original_estimated_cost=0,
                        estimated_savings=0,
                    )

        # ── Step 4: Calculate original cost ───────────────────────
        original_cost = estimate_cost(model, input_tokens, max_tokens)

        # ── Step 5: Model optimization ────────────────────────────
        recommended_model = model
        if self.config.auto_downgrade and complexity == TaskComplexity.SIMPLE:
            downgraded = DOWNGRADE_MAP.get(model)
            if downgraded and downgraded in PRICING:
                recommended_model = downgraded

        # Check blocked models
        if recommended_model in self.config.blocked_models:
            recommended_model = get_cheapest_alternative(
                recommended_model, detect_provider(recommended_model)
            )

        # ── Step 6: Calculate optimized cost ──────────────────────
        optimized_cost = estimate_cost(recommended_model, input_tokens, max_tokens)

        # ── Step 7: Budget enforcement ────────────────────────────
        allowed = True
        message_parts = []

        # Add policy enforcement message if any
        if policy_message:
            message_parts.append(policy_message)

        if max_budget is not None and optimized_cost > max_budget:
            allowed = False
            message_parts.append(
                f"Estimated cost ${optimized_cost:.6f} exceeds budget ${max_budget:.2f}"
            )

        if self.config.daily_budget_usd is not None:
            # Note: daily budget tracking would need persistent state
            # For now, just check per-request
            pass

        # ── Step 8: Build savings info ────────────────────────────
        savings = original_cost - optimized_cost
        savings_pct = (savings / original_cost * 100) if original_cost > 0 else 0

        if recommended_model != model:
            message_parts.append(
                f"Downgraded {model} -> {recommended_model} (saves {savings_pct:.0f}%)"
            )

        if not message_parts:
            message_parts.append("Request optimized")

        # ── Step 9: Build decision ────────────────────────────────
        decision = Decision(
            recommended_model=recommended_model,
            max_tokens=max_tokens,
            estimated_cost=optimized_cost,
            original_model=model,
            input_tokens=input_tokens,
            complexity=complexity,
            allowed=allowed,
            savings_pct=savings_pct,
            message="; ".join(message_parts),
            original_estimated_cost=original_cost,
            estimated_savings=savings,
        )

        # ── Step 10: Track decision (async, non-blocking) ─────────
        # Use team as project_id if specified (enables per-team spend tracking)
        self._telemetry.track_decision(
            decision.to_dict(), project_id=team or self.project_id
        )

        return decision

    def report(
        self,
        decision: Decision,
        actual_tokens: Optional[int] = None,
        output_tokens: Optional[int] = None,
        latency_ms: Optional[float] = None,
    ):
        """
        Report actual usage after the LLM call completes.

        This is optional but recommended — it feeds the analytics dashboard
        with real usage data for cost tracking and recommendations.

        Args:
            decision: The Decision object from optimize().
            actual_tokens: Total tokens from the LLM response (input + output).
            output_tokens: Output tokens only (if available).
            latency_ms: Request latency in milliseconds.
        """
        out_tokens = output_tokens or (
            (actual_tokens - decision.input_tokens)
            if actual_tokens
            else decision.max_tokens
        )
        total = actual_tokens or (decision.input_tokens + out_tokens)

        actual_cost = estimate_cost(
            decision.recommended_model, decision.input_tokens, out_tokens
        )

        usage = UsageReport(
            request_id=decision.request_id,
            model=decision.recommended_model,
            input_tokens=decision.input_tokens,
            output_tokens=out_tokens,
            total_tokens=total,
            actual_cost=actual_cost,
            latency_ms=latency_ms,
        )

        self._telemetry.track_usage(
            usage.to_dict(), project_id=self.project_id
        )

    def flush(self):
        """Force-flush any buffered telemetry events."""
        self._telemetry.flush()

    def shutdown(self):
        """Gracefully shutdown (flushes telemetry)."""
        self._telemetry.shutdown()
