"""
Phoenix Sync - Fetch trajectories from Arize Phoenix and generate guidelines.

This module provides functionality to:
1. Fetch agent trajectories from Phoenix's REST API
2. Deduplicate already-processed trajectories
3. Generate guidelines from new trajectories
4. Store both trajectories and guidelines in the Evolve backend
"""

import json
import logging
import urllib.request
import warnings
from dataclasses import dataclass
from typing import Any

from altk_evolve.config.phoenix import phoenix_settings
from altk_evolve.config.evolve import evolve_config
from altk_evolve.frontend.client.evolve_client import EvolveClient
from altk_evolve.llm.guidelines.guidelines import generate_guidelines
from altk_evolve.schema.core import Entity
from altk_evolve.schema.exceptions import NamespaceNotFoundException

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("evolve.sync.phoenix")


@dataclass
class SyncResult:
    """Result of a sync operation."""

    processed: int
    skipped: int
    guidelines_generated: int
    errors: list[str]

    @property
    def tips_generated(self) -> int:
        """Temporary compatibility alias for one release cycle."""
        warnings.warn(
            "SyncResult.tips_generated is deprecated; use SyncResult.guidelines_generated instead.",
            DeprecationWarning,
            stacklevel=2,
        )
        return self.guidelines_generated


class PhoenixSync:
    """Sync trajectories from Arize Phoenix to Evolve."""

    def __init__(
        self,
        phoenix_url: str | None = None,
        namespace_id: str | None = None,
        project: str | None = None,
    ):
        self.phoenix_url = phoenix_url or phoenix_settings.url
        self.project = project or phoenix_settings.project
        self.namespace_id = namespace_id or evolve_config.namespace_id
        self.client = EvolveClient()

    def _ensure_namespace(self):
        """Ensure the target namespace exists."""
        try:
            self.client.get_namespace_details(self.namespace_id)
        except NamespaceNotFoundException:
            self.client.create_namespace(self.namespace_id)
            logger.info(f"Created namespace: {self.namespace_id}")

    def _fetch_spans(self, limit: int = 1000) -> list[dict]:
        """Fetch spans from Phoenix, handling pagination."""
        spans: list[dict] = []
        cursor = None

        while True:
            url = f"{self.phoenix_url}/v1/projects/{self.project}/spans?limit={min(limit - len(spans), 100)}"
            if cursor:
                url += f"&cursor={cursor}"

            try:
                with urllib.request.urlopen(url, timeout=30) as response:
                    data = json.loads(response.read().decode())
            except Exception as e:
                logger.error(f"Failed to fetch spans from Phoenix: {e}")
                raise

            spans.extend(data.get("data", []))
            cursor = data.get("next_cursor")

            if not cursor or len(spans) >= limit:
                break

        return spans

    def _get_processed_span_ids(self) -> set[str]:
        """Get span_ids that have already been processed."""
        try:
            entities = self.client.search_entities(
                namespace_id=self.namespace_id,
                filters={"type": "trajectory"},
                limit=10000,
            )
            return {str(e.metadata.get("span_id")) for e in entities if e.metadata and e.metadata.get("span_id")}
        except NamespaceNotFoundException:
            return set()

    def _format_payload_summary(self, payload: Any) -> str:
        """Format a payload summary for secure logging (avoid PII)."""
        type_name = type(payload).__name__

        if isinstance(payload, str):
            length = len(payload)
            preview = payload[:50] + "..." if length > 50 else payload
            # Replace newlines in preview to keep logs on one line
            preview = preview.replace("\n", "\\n")
            return f"<{type_name} length={length} preview='{preview}'>"

        if isinstance(payload, (dict, list)):
            length = len(payload)
            return f"<{type_name} length={length}>"

        return f"<{type_name}>"

    def _parse_content(self, content: Any) -> Any:
        """Parse content which may be a string representation of a list/dict."""
        if isinstance(content, str):
            try:
                return json.loads(content)
            except json.JSONDecodeError:
                try:
                    import ast

                    return ast.literal_eval(content)
                except (ValueError, SyntaxError):
                    return content
        return content

    def _extract_messages_from_span(self, span: dict) -> list[dict]:
        """Extract messages from a single span's attributes."""
        attrs = span.get("attributes") or {}
        messages = []

        # Try OpenInference schema first (input_messages/output_messages)
        # Note: Phoenix API might return these as JSON strings or lists depending on version
        input_msgs = attrs.get("llm.input_messages")
        output_msgs = attrs.get("llm.output_messages")

        # If input_messages is missing, try parsing input.value
        if input_msgs is None:
            input_val = attrs.get("input.value")
            if input_val:
                try:
                    parsed_input = self._parse_content(input_val)
                    if isinstance(parsed_input, dict) and "messages" in parsed_input:
                        input_msgs = parsed_input["messages"]
                    elif isinstance(parsed_input, list):  # rare but possible
                        input_msgs = parsed_input
                except Exception as e:
                    logger.debug(f"Failed to parse input.value: {e}. Payload: {self._format_payload_summary(input_val)}")

        if input_msgs:
            # Handle OpenInference format
            # Ensure it's a list
            if isinstance(input_msgs, str):
                input_msgs = self._parse_content(input_msgs)

            if isinstance(input_msgs, list):
                for i, msg in enumerate(input_msgs):
                    # OpenInference often uses message.role / message.content keys in flattened export
                    # but via API it might be cleaner. Let's handle dict access safely.
                    if isinstance(msg, str):
                        try:
                            msg = self._parse_content(msg)
                        except Exception as e:
                            logger.debug(f"Failed to parse input message string: {e}. Payload: {self._format_payload_summary(msg)}")

                    if not isinstance(msg, dict):
                        continue
                    role = msg.get("message.role") or msg.get("role")
                    content = msg.get("message.content") or msg.get("content")
                    tool_calls = msg.get("message.tool_calls") or msg.get("tool_calls")

                    if role:
                        mapped_msg = {
                            "index": i,
                            "type": "prompt",
                            "role": role,
                            "content": self._parse_content(content),
                        }
                        if tool_calls:
                            mapped_msg["tool_calls"] = tool_calls
                        messages.append(mapped_msg)

        # Handle Output/Completion from OpenInference
        if output_msgs is None:
            output_val = attrs.get("output.value")
            if output_val:
                try:
                    parsed_output = self._parse_content(output_val)
                    # output.value is often just the string content or a list of choices
                    if isinstance(parsed_output, list) and len(parsed_output) > 0 and "message" in parsed_output[0]:
                        output_msgs = [c["message"] for c in parsed_output]
                    elif isinstance(parsed_output, dict) and "choices" in parsed_output:  # OpenAI response format
                        output_msgs = [c["message"] for c in parsed_output["choices"]]
                    else:
                        # Fallback for simple string output
                        output_msgs = [{"role": "assistant", "content": output_val}]
                        # pass
                except Exception as e:
                    logger.debug(f"Failed to parse output.value: {e}. Payload: {self._format_payload_summary(output_val)}")

        if output_msgs:
            if isinstance(output_msgs, str):
                output_msgs = self._parse_content(output_msgs)

            if isinstance(output_msgs, list):
                for i, msg in enumerate(output_msgs):
                    if isinstance(msg, str):
                        try:
                            msg = self._parse_content(msg)
                        except Exception as e:
                            logger.debug(f"Failed to parse output message string: {e}. Payload: {self._format_payload_summary(msg)}")

                    if not isinstance(msg, dict):
                        continue
                    role = msg.get("message.role") or msg.get("role")
                    content = msg.get("message.content") or msg.get("content")
                    tool_calls = msg.get("message.tool_calls") or msg.get("tool_calls")

                    if role:
                        mapped_msg = {
                            "index": i,
                            "type": "completion",
                            "role": role,
                            "content": self._parse_content(content),
                        }
                        if tool_calls:
                            mapped_msg["tool_calls"] = tool_calls
                        messages.append(mapped_msg)

        if messages:
            return messages

        # Fallback to GenAI semantic conventions (original code)
        # Extract prompt messages
        prompt_indices = set()
        for key in attrs:
            if key.startswith("gen_ai.prompt.") and key.endswith(".role"):
                idx = int(key.split(".")[2])
                prompt_indices.add(idx)

        for i in sorted(prompt_indices):
            role = attrs.get(f"gen_ai.prompt.{i}.role")
            content = attrs.get(f"gen_ai.prompt.{i}.content")
            if role and content is not None:
                messages.append(
                    {
                        "index": i,
                        "type": "prompt",
                        "role": role,
                        "content": self._parse_content(content),
                    }
                )

        # Extract completion messages
        completion_indices = set()
        for key in attrs:
            if key.startswith("gen_ai.completion.") and key.endswith(".role"):
                idx = int(key.split(".")[2])
                completion_indices.add(idx)

        for i in sorted(completion_indices):
            role = attrs.get(f"gen_ai.completion.{i}.role")
            content = attrs.get(f"gen_ai.completion.{i}.content")
            if role and content is not None:
                messages.append(
                    {
                        "index": i,
                        "type": "completion",
                        "role": role,
                        "content": self._parse_content(content),
                    }
                )

        return messages

    def _convert_to_openai_format(self, content: Any, role: str) -> dict:
        """Convert Anthropic message format to OpenAI format."""
        if isinstance(content, str):
            return {"role": role, "content": content}

        if not isinstance(content, list):
            return {"role": role, "content": str(content)}

        text_parts = []
        tool_calls = []
        tool_results = []
        thinking_parts = []

        for block in content:
            if not isinstance(block, dict):
                text_parts.append(str(block))
                continue

            block_type = block.get("type")

            if block_type == "text":
                text = block.get("text", "")
                if text and text != "(no content)":
                    text_parts.append(text)

            elif block_type == "thinking":
                thinking = block.get("thinking", "")
                if thinking:
                    thinking_parts.append(thinking)

            elif block_type == "tool_use":
                tool_calls.append(
                    {
                        "id": block.get("id", ""),
                        "type": "function",
                        "function": {
                            "name": block.get("name", ""),
                            "arguments": json.dumps(block.get("input", {})),
                        },
                    }
                )

            elif block_type == "tool_result":
                tool_results.append(
                    {
                        "tool_call_id": block.get("tool_use_id", ""),
                        "content": block.get("content", ""),
                        "is_error": block.get("is_error", False),
                    }
                )

        if role == "assistant":
            msg: dict[str, str | list | None] = {"role": "assistant"}
            if thinking_parts:
                msg["thinking"] = "\n\n".join(thinking_parts)
            if text_parts:
                msg["content"] = "\n\n".join(text_parts)
            elif not tool_calls:
                msg["content"] = None
            if tool_calls:
                msg["tool_calls"] = tool_calls
            return msg

        elif role == "user" and tool_results:
            return {"role": "tool", "tool_results": tool_results}

        else:
            content_text = "\n\n".join(text_parts) if text_parts else ""
            return {"role": role, "content": content_text}

    def _extract_trajectory(self, span: dict) -> dict:
        """Extract a complete trajectory from a span."""
        attrs = span.get("attributes") or {}
        messages = self._extract_messages_from_span(span)

        openai_messages = []

        for msg in messages:
            role = msg["role"]
            content = msg["content"]
            converted = self._convert_to_openai_format(content, role)

            if converted.get("role") == "tool" and "tool_results" in converted:
                for result in converted["tool_results"]:
                    openai_messages.append(
                        {
                            "role": "tool",
                            "tool_call_id": result["tool_call_id"],
                            "content": result["content"],
                        }
                    )
            else:
                openai_messages.append(converted)

        return {
            "trace_id": span["context"]["trace_id"],
            "span_id": span["context"]["span_id"],
            "model": attrs.get("gen_ai.request.model", "unknown"),
            "timestamp": span.get("start_time"),
            "messages": openai_messages,
            "usage": {
                "prompt_tokens": next(
                    (
                        v
                        for v in [
                            attrs.get("gen_ai.usage.prompt_tokens"),
                            attrs.get("llm.token_count.prompt"),
                            attrs.get("llm.usage.prompt_tokens"),
                        ]
                        if v is not None
                    ),
                    None,
                ),
                "completion_tokens": next(
                    (
                        v
                        for v in [
                            attrs.get("gen_ai.usage.completion_tokens"),
                            attrs.get("llm.token_count.completion"),
                            attrs.get("llm.usage.completion_tokens"),
                        ]
                        if v is not None
                    ),
                    None,
                ),
                "total_tokens": next(
                    (
                        v
                        for v in [
                            attrs.get("gen_ai.usage.total_tokens"),
                            attrs.get("llm.token_count.total"),
                            attrs.get("llm.usage.total_tokens"),
                        ]
                        if v is not None
                    ),
                    None,
                ),
            },
        }

    def _clean_trajectory(self, trajectory: dict) -> dict:
        """Clean up a trajectory by removing system reminders."""
        import re

        cleaned_messages = []

        for msg in trajectory.get("messages", []):
            if not msg.get("content") and not msg.get("tool_calls"):
                continue

            if msg.get("content"):
                content = msg["content"]
                if isinstance(content, str):
                    content = re.sub(
                        r"<system-reminder>.*?</system-reminder>",
                        "",
                        content,
                        flags=re.DOTALL,
                    ).strip()
                    if not content:
                        continue
                    msg = {**msg, "content": content}

            cleaned_messages.append(msg)

        return {**trajectory, "messages": cleaned_messages}

    def _process_trajectory(self, trajectory: dict) -> int:
        """Process a single trajectory: store it and generate guidelines.

        Returns the number of guidelines generated.
        """
        # Store trajectory as a single entity with all messages
        messages = trajectory.get("messages", [])
        if messages:
            entity = Entity(
                type="trajectory",
                content=json.dumps(messages),
                metadata={
                    "trace_id": trajectory["trace_id"],
                    "span_id": trajectory["span_id"],
                    "model": trajectory["model"],
                    "timestamp": trajectory["timestamp"],
                    "message_count": len(messages),
                    "usage": trajectory.get("usage"),
                },
            )

            self.client.update_entities(
                namespace_id=self.namespace_id,
                entities=[entity],
                enable_conflict_resolution=False,
            )

        # Generate guidelines from the trajectory (returns one result per subtask)
        results = generate_guidelines(trajectory["messages"])

        guideline_entities = [
            Entity(
                type="guideline",
                content=guideline.content,
                metadata={
                    "category": guideline.category,
                    "rationale": guideline.rationale,
                    "trigger": guideline.trigger,
                    "implementation_steps": guideline.implementation_steps,
                    "source_task_id": trajectory["trace_id"],
                    "source_span_id": trajectory["span_id"],
                    "task_description": result.task_description,
                    "creation_mode": "auto-phoenix",
                },
            )
            for result in results
            for guideline in result.guidelines
        ]
        if guideline_entities:
            self.client.update_entities(
                namespace_id=self.namespace_id,
                entities=guideline_entities,
                enable_conflict_resolution=True,
            )

        return sum(len(r.guidelines) for r in results)

    def sync(
        self,
        limit: int = 100,
        include_errors: bool = False,
    ) -> SyncResult:
        """
        Fetch new trajectories from Phoenix and generate guidelines.

        Args:
            limit: Maximum number of spans to fetch from Phoenix
            include_errors: Whether to include failed/error spans

        Returns:
            SyncResult with counts of processed, skipped, and guidelines generated
        """
        logger.info(f"Starting sync from {self.phoenix_url} to namespace '{self.namespace_id}'")

        self._ensure_namespace()

        # Fetch spans from Phoenix
        spans = self._fetch_spans(limit)
        logger.info(f"Fetched {len(spans)} spans from Phoenix")

        # Get already processed span IDs
        processed_ids = self._get_processed_span_ids()
        logger.info(f"Found {len(processed_ids)} already processed spans")

        processed = 0
        skipped = 0
        guidelines_generated = 0
        errors = []

        for span in spans:
            # Filter to LLM request spans - accept any span with prompt attributes
            # if span.get("name") != "litellm_request":
            #     continue

            # Filter errors if requested
            if not include_errors and span.get("status_code") == "ERROR":
                continue

            # Check if already processed
            span_id = span.get("context", {}).get("span_id")
            if span_id in processed_ids:
                skipped += 1
                continue

            # Only include spans with actual messages or GenAI/LLM prompt attributes
            attrs = span.get("attributes") or {}
            has_gen_ai = any(k.startswith("gen_ai.prompt.") for k in attrs)
            has_llm_msgs = "llm.input_messages" in attrs or "input.value" in attrs

            if not (has_gen_ai or has_llm_msgs):
                continue

            try:
                trajectory = self._extract_trajectory(span)
                trajectory = self._clean_trajectory(trajectory)

                if trajectory["messages"]:
                    guidelines_count = self._process_trajectory(trajectory)
                    processed += 1
                    guidelines_generated += guidelines_count
                    logger.info(f"Processed span {span_id[:12]}... - generated {guidelines_count} guidelines")
            except Exception as e:
                error_msg = f"Error processing span {span_id}: {e}"
                logger.exception(error_msg)
                errors.append(error_msg)

        result = SyncResult(
            processed=processed,
            skipped=skipped,
            guidelines_generated=guidelines_generated,
            errors=errors,
        )

        logger.info(
            f"Sync complete: {processed} processed, {skipped} skipped, {guidelines_generated} guidelines generated, {len(errors)} errors"
        )

        return result
