"""
vgovern SDK - Tracer

Manages trace lifecycle and provides context for tracking agent executions.
"""

import logging
from contextlib import contextmanager
from dataclasses import dataclass, field
from datetime import datetime,timezone
from typing import Any, Dict, Generator, Optional, List, TYPE_CHECKING
import contextvars
from uuid import UUID, uuid4
from threading import Lock


from .sdk_types import (
    SDKConfig,
    StepData,
    StepType,
    TraceInfo,
    TraceMetadata,
    TraceStatus,
    GovernanceResponse,
    GovernanceDecision,
    TokenUsage,
    StepNode,
    Trace
    
)
from .client import GovernanceClient
from .utils import TraceLogger, get_logger, generate_step_id, redact_sensitive_data

logger = get_logger(__name__)

from threading import Lock
from contextvars import ContextVar

_trace_context_stack: ContextVar[list] = ContextVar("_trace_context_stack", default=[])

@dataclass
class ActiveTrace:
    """Holds state for an active trace."""
    info: TraceInfo
    steps: List[StepData] = field(default_factory=list)
    metadata: Dict[str, Any] = field(default_factory=dict)
    lock: Lock = field(default_factory=Lock)

    def add_step(self, step: StepData) -> None:
        with self.lock:
            self.steps.append(step)

class Tracer:
    """
    Manages trace lifecycle and provides context for tracking agent executions.
    Thread-safe implementation with fail-safe guarantees.
    """

    def __init__(self, config: SDKConfig):
        self.config = config
        self._active_traces: Dict[UUID, ActiveTrace] = {}
        self._lock = Lock()
        self._logger = TraceLogger(logger)
        self._client = GovernanceClient(config)

    def start_trace(self, metadata: Optional[TraceMetadata] = None) -> UUID:
        trace_id = uuid4()
        now = datetime.now(timezone.utc)
        trace_metadata = metadata or TraceMetadata(agent_id=self.config.agent_id)
        trace_info = TraceInfo(
            trace_id=trace_id,
            agent_id=trace_metadata.agent_id,
            start_time=now,
            status=TraceStatus.RUNNING,
            session_id=trace_metadata.session_id,
            metadata=trace_metadata.tags,
        )
        active_trace = ActiveTrace(info=trace_info)
        with self._lock:
            self._active_traces[trace_id] = active_trace
        self._logger = self._logger.with_trace(trace_id)
        self._logger.info(f"Started trace: {trace_id}")
        # Initialize trace tree
        trace = Trace.from_trace_info(trace_info)
        trace.user_input = trace_metadata.tags.get("user_input") if trace_metadata.tags else None
        active_trace.trace_tree = trace
        try:
            self._client.send_trace_start(
                trace_info=trace_info,
                metadata=trace_metadata.tags if trace_metadata.tags else {},
            )
        except Exception as e:
            self._logger.warning(f"Failed to send trace start: {e}")
        return trace_id

    def end_trace(
        self,
        trace_id: UUID,
        status: TraceStatus = TraceStatus.COMPLETED,
        metadata: Optional[Dict[str, Any]] = None,
    ) -> None:
        with self._lock:
            active_trace = self._active_traces.pop(trace_id, None)
        if active_trace is None:
            self._logger.warning(f"Trace not found: {trace_id}")
            return
        object.__setattr__(active_trace.info, "status", status)
        end_time = datetime.now(timezone.utc)
        self._logger.info(f"Ended trace: {trace_id} with status: {status.value}")
        # Update trace tree end time
        if active_trace.trace_tree:
            active_trace.trace_tree.end_time = end_time
            # Print tree if console trace is enabled
            if self.config.enable_console_trace:
                Tracer.print_trace_tree(active_trace.trace_tree)
        try:
            self._client.send_trace_end(
                trace_id=trace_id,
                status=status,
                end_time=end_time,
            )
        except Exception as e:
            self._logger.warning(f"Failed to send trace end: {e}")

    def record_step(
        self,
        trace_id: UUID,
        step_type: StepType,
        model: Optional[str] = None,
        prompt: Optional[str] = None,
        response: Optional[str] = None,
        token_usage: Optional[Dict[str, int]] = None,
        latency_ms: Optional[int] = None,
        estimated_cost_usd: Optional[float] = None,
        metadata: Optional[Dict[str, Any]] = None,
    ) -> str:
        step_id = generate_step_id()
        now = datetime.now(timezone.utc)
        if self.config.redact_sensitive_data:
            prompt = redact_sensitive_data(prompt) if prompt else None
            response = redact_sensitive_data(response) if response else None
        token_usage_obj = None
        if token_usage:
            token_usage_obj = TokenUsage(
                prompt_tokens=token_usage.get("prompt_tokens", 0),
                completion_tokens=token_usage.get("completion_tokens", 0),
                total_tokens=token_usage.get("total_tokens", 0),
            )
        step_data = StepData(
            step_id=step_id,
            trace_id=trace_id,
            timestamp=now,
            step_type=step_type,
            model=model,
            prompt=prompt,
            response=response,
            token_usage=token_usage_obj,
            latency_ms=latency_ms,
            estimated_cost_usd=estimated_cost_usd,
            metadata=metadata or {},
        )
        with self._lock:
            active_trace = self._active_traces.get(trace_id)
            if active_trace:
                active_trace.add_step(step_data)
        self._logger.info(f"Recorded step: {step_id} (type: {step_type.value})")
        if self.config.enable_async_reporting:
            try:
                import threading
                def _send_async():
                    self._client.send_step(step_data)
                thread = threading.Thread(target=_send_async, daemon=True)
                thread.start()
            except Exception as e:
                self._logger.warning(f"Failed to send step async: {e}")
        else:
            try:
                self._client.send_step(step_data)
            except Exception as e:
                self._logger.warning(f"Failed to send step: {e}")
        return step_id

    def check_policy(
        self,
        trace_id: UUID,
        step_type: StepType,
        model: Optional[str] = None,
        token_usage: Optional[Dict[str, int]] = None,
        estimated_cost_usd: Optional[float] = None,
        metadata: Optional[Dict[str, Any]] = None,
    ) -> GovernanceResponse:
        token_usage_obj = None
        if token_usage:
            token_usage_obj = TokenUsage(
                prompt_tokens=token_usage.get("prompt_tokens", 0),
                completion_tokens=token_usage.get("completion_tokens", 0),
            )
        step_data = StepData(
            step_id="policy_check",
            trace_id=trace_id,
            timestamp=datetime.now(timezone.utc),
            step_type=step_type,
            model=model,
            token_usage=token_usage_obj,
            estimated_cost_usd=estimated_cost_usd,
            metadata=metadata or {},
        )
        response = self._client.check_policy(trace_id, step_data)
        self._logger.info(
            f"Policy check: {response.decision.value} (policy_id: {response.policy_id or 'N/A'})"
        )
        return response

    @contextmanager
    def trace_context(
        self,
        metadata: Optional[TraceMetadata] = None,
    ) -> Generator[UUID, None, None]:
        trace_id = self.start_trace(metadata)
        try:
            yield trace_id
        except Exception as e:
            from .sdk_types import GovernanceHaltError
            if isinstance(e, GovernanceHaltError):
                self.end_trace(trace_id, TraceStatus.HALTED)
            else:
                self.end_trace(trace_id, TraceStatus.ERROR)
            raise
        else:
            self.end_trace(trace_id, TraceStatus.COMPLETED)

    def get_trace_info(self, trace_id: UUID) -> Optional[TraceInfo]:
        with self._lock:
            active_trace = self._active_traces.get(trace_id)
            if active_trace:
                return active_trace.info
        return None

    def cleanup(self) -> None:
        try:
            self._client.close()
        except Exception as e:
            self._logger.warning(f"Failed to close client: {e}")


    def start_step(
        self,
        trace_id: UUID,
        name: str,
        step_type: StepType,
        input_data: Optional[Dict[str, Any]] = None,
    ) -> StepNode:
        """Start a new execution step and add it to the trace tree."""
        step_id = generate_step_id()
        now = datetime.now(timezone.utc)
        
        # Get parent_id from context stack
        stack = _trace_context_stack.get([])
        parent_id = stack[-1] if stack else None
        
        step_node = StepNode(
            step_id=step_id,
            trace_id=trace_id,
            name=name,
            step_type=step_type,
            parent_id=parent_id,
            start_time=now,
            input_data=input_data,
        )
        
        # Add to trace tree
        with self._lock:
            active_trace = self._active_traces.get(trace_id)
            if active_trace:
                if parent_id:
                    # Find parent and add child
                    def add_to_parent(node):
                        for child in node.children:
                            if child.step_id == parent_id:
                                child.add_child(step_node)
                                return True
                            if add_to_parent(child):
                                return True
                        return False
                    
                    # Search in root steps first
                    found = False
                    for root_step in (active_trace.trace_tree.root_steps if active_trace.trace_tree else []):
                        if add_to_parent(root_step):
                            found = True
                            break
                else:
                    # Add as root step
                    if active_trace.trace_tree:
                        active_trace.trace_tree.add_root_step(step_node)
        
        # Push step_id to context stack
        new_stack = stack + [step_id]
        _trace_context_stack.set(new_stack)
        
        self._logger.info(f"Started step: {step_id} (name: {name}, type: {step_type.value})")
        return step_node


    def end_step(
        self,
        step_id: str,
        output_data: Optional[Dict[str, Any]] = None,
    ) -> None:
        """End an execution step and compute duration."""
        now = datetime.now(timezone.utc)
        
        # Pop from context stack
        stack = _trace_context_stack.get([])
        if stack and stack[-1] == step_id:
            new_stack = stack[:-1]
            _trace_context_stack.set(new_stack)
        
        # Find and update step
        def find_and_update(node):
            if node.step_id == step_id:
                node.end_time = now
                if node.start_time:
                    delta = now - node.start_time
                    node.duration_ms = int(delta.total_seconds() * 1000)
                node.output_data = output_data
                return True
            for child in node.children:
                if find_and_update(child):
                    return True
            return False
        
        # Search in all active traces
        with self._lock:
            for active_trace in self._active_traces.values():
                if active_trace.trace_tree:
                    for root_step in active_trace.trace_tree.root_steps:
                        find_and_update(root_step)
        
        self._logger.info(f"Ended step: {step_id}")


    def add_policy_hit(
        self,
        step_id: str,
        policy: str,
        decision: str,
        reason: str,
    ) -> None:
        """Attach a policy hit result to a step.
        
        This method is fail-safe - errors will not interrupt agent execution.
        
        Args:
            step_id: The ID of the step to attach the policy hit to
            policy: Name of the policy that was evaluated
            decision: Policy decision (ALLOW, HALT, WARN)
            reason: Human-readable explanation of the decision
        """
        try:
            policy_hit = {
                "policy": policy,
                "decision": decision,
                "reason": reason,
            }
            
            # Find the step node by step_id
            def find_and_attach(node):
                if node.step_id == step_id:
                    if node.policy_hits is None:
                        node.policy_hits = []
                    node.policy_hits.append(policy_hit)
                    return True
                for child in node.children:
                    if find_and_attach(child):
                        return True
                return False
            
            # Search in all active traces
            found = False
            with self._lock:
                for active_trace in self._active_traces.values():
                    if active_trace.trace_tree:
                        for root_step in active_trace.trace_tree.root_steps:
                            if find_and_attach(root_step):
                                found = True
                                break
                        if found:
                            break
            
            if found:
                self._logger.info(f"Added policy hit to step {step_id}: {policy} -> {decision}")
            else:
                self._logger.warning(f"Step not found for policy hit: {step_id}")
        except Exception as e:
            # Fail-safe: log error but don't interrupt agent execution
            self._logger.warning(f"Failed to add policy hit: {e}")
    @contextmanager
    def step_context(
        self,
        trace_id: UUID,
        name: str,
        step_type: StepType,
        input_data: Optional[Dict[str, Any]] = None,
    ) -> Generator[StepNode, None, None]:
        """Context manager for automatic step lifecycle management."""
        step_node = self.start_step(trace_id, name, step_type, input_data)
        try:
            yield step_node
        except Exception:
            self.end_step(step_node.step_id)
            raise
        else:
            self.end_step(step_node.step_id)


    @staticmethod
    def print_trace_tree(trace: Trace) -> None:
          """Print trace tree in Console_trace.md format without color highlighting."""
          if not trace:
              print("[Trace] No trace data available")
              return

          # Helper to get step emoji based on type
          def get_step_emoji(step_type: StepType) -> str:
              if step_type == StepType.LLM_CALL:
                  return "🧠"
              elif step_type == StepType.TOOL_CALL:
                  return "🛠"
              elif step_type == StepType.FINAL_RESPONSE:
                  return "✅"
              return "📋"

          # Helper to check if node has HALT decision
          def has_halt_decision(node: StepNode) -> bool:
              if not node.policy_hits:
                  return False
              return any(hit.get("decision", "").upper() == "HALT" for hit in node.policy_hits)

          # Check if trace has any HALT decisions
          trace_halted = any(has_halt_decision(root) for root in trace.root_steps)

          # Find halt reason and step if halted
          halt_reason = ""
          halt_step = ""
          if trace_halted:
              for root in trace.root_steps:
                  if has_halt_decision(root):
                      for hit in root.policy_hits:
                          if hit.get("decision", "").upper() == "HALT":
                              halt_reason = hit.get("reason", "No reason")
                              halt_step = root.name
                              break
                      break

          # Print trace header
          print(f"🧩 Trace: {str(trace.trace_id)[:8]}")
          print()

          # Print HALT summary at top if trace was halted
          if trace_halted:
              print(f"⛔ TRACE HALTED")
              print()
              print(f"├─ 🧾 Reason:")
              print(f"│  {halt_reason}")
              print(f"├─ 🛑 Halted At:")
              print(f"│  Step: {halt_step}")
              print(f"├─ 📊 Accumulated Cost:")
              print(f"│  $0.0000")
              print(f"└─ ⏱ Elapsed Time:")
              print(f"    {trace.get_total_duration_ms() / 1000:.2f}s")
              print()

          def _format_dict(data: Any) -> str:
              """Format dict or other data for display."""
              if isinstance(data, dict):
                  try:
                      import json
                      return json.dumps(data, ensure_ascii=False)
                  except Exception:
                      pass
              return str(data)

          def _print_node(node: StepNode, prefix: str = "", is_last: bool = False) -> None:
              # Print current node header
              connector = "└─" if is_last else "├─"
              duration = node.get_duration_display()
              emoji = get_step_emoji(node.step_type)

              # Determine step type label
              step_type_label = ""
              if node.step_type == StepType.LLM_CALL:
                  step_type_label = "(LLM)"
              elif node.step_type == StepType.TOOL_CALL:
                  step_type_label = "(工具调用)"
              elif node.step_type == StepType.FINAL_RESPONSE:
                  step_type_label = "(最终响应)"

              print(f"{prefix}{connector} {emoji} Step: {node.name} {step_type_label}")
              print(f"{prefix}│  ├─ ⏱ Duration: {duration}")

              # Print input data
              if node.input_data:
                  if node.step_type == StepType.TOOL_CALL:
                      print(f"{prefix}│  ├─ 📥 Parameters:")
                      print(f"{prefix}│  │  {_format_dict(node.input_data)}")
                  elif node.step_type == StepType.LLM_CALL:
                      print(f"{prefix}│  ├─ 📥 Input:")
                      print(f"{prefix}│  │  {_format_dict(node.input_data)}")

              # Print output data
              if node.output_data:
                  if node.step_type == StepType.TOOL_CALL:
                      result = node.output_data.get("result")
                      if result:
                          print(f"{prefix}│  ├─ 📤 Result:")
                          print(f"{prefix}│  │  {_format_dict(result)}")
                  elif node.step_type == StepType.FINAL_RESPONSE:
                      answer = node.output_data.get("answer")
                      if answer:
                          print(f"{prefix}│  ├─ 📤 Output:")
                          print(f"{prefix}│  │  {answer}")

              # Print policy hits
              if node.policy_hits:
                  for hit in node.policy_hits:
                      decision = hit.get("decision", "unknown").upper()
                      policy_name = hit.get("policy", "unknown")

                      if decision == "HALT":
                          print(f"{prefix}│  └─ ⛔ POLICY HALT")
                          print(f"{prefix}│     ├─ Rule: {policy_name}")
                          print(f"{prefix}│     ├─ Limit: $0.05")
                          print(f"{prefix}│     ├─ Current: $0.000")
                          print(f"{prefix}│     └─ Action: Execution stopped immediately")
                          break
                      else:
                          print(f"{prefix}│  └─ ⚠️ Policy Hit:")
                          print(f"{prefix}│     ├─ Rule: {policy_name}")
                          print(f"{prefix}│     ├─ Limit: $0.05")
                          print(f"{prefix}│     └─ Current: $0.000")

              # Print children
              child_prefix = prefix + ("   " if is_last else "│  ")
              for i, child in enumerate(node.children):
                  _print_node(child, child_prefix, i == len(node.children) - 1)

          # Print all root steps
          for i, root_step in enumerate(trace.root_steps):
              _print_node(root_step, "", i == len(trace.root_steps) - 1)

          # Print trace footer
          print()
          status = "HALTED" if trace_halted else "COMPLETED"
          status_emoji = "⛔" if trace_halted else "✅"
          print(f"└─ {status_emoji} Trace Status: {status}")
          print(f"   ├─ Total Duration: {trace.get_total_duration_ms() / 1000:.2f}s")
          print(f"   ├─ 📊 Total Cost: $0.0000")
          print(f"   └─ 🧾 Steps: {len(trace.root_steps)}")

          # Print HALT summary at bottom if trace was halted
          if trace_halted:
              print()
              print(f"⛔ TRACE HALTED")
              print()
              print(f"├─ 🧾 Reason:")
              print(f"│  {halt_reason}")
              print(f"├─ 🛑 Halted At:")
              print(f"│  Step: {halt_step}")
              print(f"├─ 📊 Accumulated Cost:")
              print(f"│  $0.0000")
              print(f"└─ ⏱ Elapsed Time:")
              print(f"    {trace.get_total_duration_ms() / 1000:.2f}s")