"""
Attribution orchestrator for releaseops.

Coordinates all analyzers and produces final attribution results
explaining why an agent took a specific action.
"""

from __future__ import annotations

import json
import logging
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

from llmhq_releaseops.attribution.analyzers.model_analyzer import ModelConfigAnalyzer
from llmhq_releaseops.attribution.analyzers.policy_analyzer import PolicyAnalyzer
from llmhq_releaseops.attribution.analyzers.prompt_analyzer import PromptAnalyzer
from llmhq_releaseops.attribution.models import Attribution, Influence
from llmhq_releaseops.storage.git_store import GitStore

logger = logging.getLogger(__name__)


class AttributionAnalyzer:
    """
    Main attribution orchestrator.

    Coordinates prompt, policy, and model analyzers to explain
    why an agent took a specific action.
    """

    def __init__(self, store: GitStore, prompt_bridge=None):
        self._store = store
        self._prompt_bridge = prompt_bridge
        self._prompt_analyzer = PromptAnalyzer()
        self._policy_analyzer = PolicyAnalyzer()
        self._model_analyzer = ModelConfigAnalyzer()

    def analyze(
        self,
        trace_data: Dict[str, Any],
        bundle_id: str,
        bundle_version: str,
    ) -> Attribution:
        """
        Analyze a trace to explain agent behavior.

        Args:
            trace_data: Trace data dict (flexible format).
            bundle_id: Bundle that produced the behavior.
            bundle_version: Version of the bundle.

        Returns:
            Attribution with influences and assessment.
        """
        bundle = self._store.load_bundle(bundle_id, bundle_version)
        action, action_type = self._extract_action(trace_data)

        all_influences: List[Influence] = []

        # Run prompt analyzer on all prompts
        for role, ref in bundle.prompts.items():
            prompt_content = self._load_prompt_content(ref.ref)
            if prompt_content:
                try:
                    influences = self._prompt_analyzer.analyze(
                        action, prompt_content,
                        prompt_ref=ref.ref,
                        trace_context=trace_data,
                    )
                    all_influences.extend(influences)
                except Exception as e:
                    logger.warning(f"Prompt analyzer failed for '{role}': {e}")

        # Run policy analyzer on all policies
        for role, ref in bundle.policies.items():
            policy_content = self._load_policy_content(ref.ref)
            if policy_content is not None:
                policy_type = self._infer_policy_type(role)
                try:
                    influences = self._policy_analyzer.analyze(
                        action, policy_type, policy_content,
                        policy_ref=ref.ref,
                        trace_context=trace_data,
                    )
                    all_influences.extend(influences)
                except Exception as e:
                    logger.warning(f"Policy analyzer failed for '{role}': {e}")

        # Run model config analyzer
        if bundle.model_config:
            try:
                influences = self._model_analyzer.analyze(
                    action, bundle.model_config.to_dict(),
                    trace_context=trace_data,
                )
                all_influences.extend(influences)
            except Exception as e:
                logger.warning(f"Model analyzer failed: {e}")

        # Sort by confidence descending
        all_influences.sort(key=lambda i: i.confidence, reverse=True)

        assessment = self._assess_overall(all_influences, action)

        return Attribution(
            trace_id=trace_data.get("trace_id", trace_data.get("id", "unknown")),
            bundle_id=bundle_id,
            bundle_version=bundle_version,
            action=action,
            action_type=action_type,
            influences=all_influences,
            overall_assessment=assessment,
            analyzed_at=datetime.now(timezone.utc).isoformat(),
        )

    def analyze_from_file(
        self,
        trace_file: Path,
        bundle_id: str,
        bundle_version: str,
    ) -> Attribution:
        """Load trace from JSON file, then analyze."""
        with open(trace_file) as f:
            trace_data = json.load(f)
        return self.analyze(trace_data, bundle_id, bundle_version)

    def _extract_action(self, trace_data: Dict[str, Any]) -> Tuple[str, str]:
        """
        Extract the action from a trace.

        Returns (action_string, action_type).
        """
        # Look for tool calls
        tool_calls = trace_data.get("tool_calls", [])
        if tool_calls:
            tc = tool_calls[-1] if isinstance(tool_calls, list) else tool_calls
            if isinstance(tc, dict):
                name = tc.get("name", tc.get("function", ""))
                args = tc.get("arguments", tc.get("args", {}))
                return f"{name}({args})" if args else name, "tool_call"

        # Look for output/response
        output = trace_data.get("output", trace_data.get("response", ""))
        if output:
            if isinstance(output, str):
                return output[:200], "response"
            return str(output)[:200], "response"

        # Look for action field
        action = trace_data.get("action", "")
        if action:
            return str(action), "decision"

        return "unknown action", "decision"

    def _assess_overall(self, influences: List[Influence], action: str) -> str:
        """Determine overall assessment based on all influences."""
        if not influences:
            return "Unclear"

        # Check for contradictions — any negated influence means the action
        # violates an explicit prohibition in the artifacts
        negated = [i for i in influences if i.is_negated and i.confidence >= 0.5]
        if negated:
            return "Contradicts artifacts"

        high_confidence = [i for i in influences if i.confidence >= 0.8]
        if high_confidence:
            return "Expected per artifacts"

        medium_confidence = [i for i in influences if i.confidence >= 0.5]
        if medium_confidence:
            return "Likely expected per artifacts"

        return "Unclear"

    def _load_prompt_content(self, ref: str) -> Optional[str]:
        """Load prompt content, returning None on failure."""
        if self._prompt_bridge:
            try:
                prompt_id = ref.split("/")[-1].split("@")[0] if "/" in ref else ref
                version = ref.split("@")[-1] if "@" in ref else None
                return self._prompt_bridge.resolve_prompt(prompt_id, version)
            except Exception as e:
                logger.warning(f"Could not load prompt '{ref}': {e}")
        return None

    def _load_policy_content(self, ref: str) -> Optional[Dict[str, Any]]:
        """Load and parse policy YAML, returning None on failure."""
        try:
            import yaml
            policy_path = self._store.repo_path / ref
            if policy_path.exists():
                with open(policy_path) as f:
                    return yaml.safe_load(f)
        except Exception as e:
            logger.warning(f"Could not load policy '{ref}': {e}")
        return None

    def _infer_policy_type(self, role: str) -> str:
        """Infer policy type from role name."""
        role_lower = role.lower()
        if "tool" in role_lower:
            return "tools"
        elif "safety" in role_lower or "guard" in role_lower:
            return "safety"
        elif "context" in role_lower or "retrieval" in role_lower:
            return "context"
        return "tools"  # Default
