import asyncio
import json
import logging
from dataclasses import dataclass
from datetime import datetime
from enum import Enum
from typing import Any, AsyncIterator, Dict, List, Optional, Tuple

import aiohttp
from dataiku.llm.python import BaseLLM
from dataiku.llm.python.types import CompletionSettings, SingleCompletionQuery, StreamCompletionResponse
from dataiku.llm.tracing import SpanBuilder
from dataikuapi.dss.llm import DSSLLMStreamedCompletionChunk, DSSLLMStreamedCompletionFooter

from connect_to_snowflake import SnowflakeConnectionHelper
from snowflake_tools_utils import make_dates_json_safe

logger = logging.getLogger("snowflakecortexagentplugin")

# Type definitions

class EventType(Enum):
    """SSE event types from Snowflake Cortex Agent."""
    TEXT_DELTA = "response.text.delta"
    THINKING_DELTA = "response.thinking.delta"
    TOOL_RESULT = "response.tool_result"
    RESPONSE = "response"


class ToolType(Enum):
    """Tool types supported by the agent."""
    CORTEX_ANALYST = "cortex_analyst_text_to_sql"
    CORTEX_SEARCH = "cortex_search"


@dataclass
class AgentConfig:
    """Configuration for Snowflake Cortex Agent."""
    connection_name: str
    agent_name: str
    database: str
    schema: str
    api_timeout: int = 300
    max_retries: int = 3
    retry_delay: int = 1
    debug: bool = False

    @classmethod
    def from_dict(cls, config: Dict[str, Any]) -> 'AgentConfig':
        """Create AgentConfig from dictionary."""
        required = ["snowflake_connection", "agent_name",
                    "agent_database", "agent_schema"]
        missing = [p for p in required if not config.get(p)]
        if missing:
            raise ValueError(
                f"Missing required parameters: {', '.join(missing)}")

        return cls(
            connection_name=config["snowflake_connection"],
            agent_name=config["agent_name"],
            database=config["agent_database"],
            schema=config["agent_schema"],
            api_timeout=config.get("api_timeout", 300),
            max_retries=config.get("max_retries", 3),
            retry_delay=config.get("retry_delay", 1),
            debug=config.get("debug", False)
        )


class SSEParser:
    """Parser for Server-Sent Events (SSE) streams."""

    CHUNK_SIZE = 64 * 1024  # 64KB chunks

    @staticmethod
    async def parse_line(line: bytes) -> Tuple[Optional[str], Optional[str]]:
        """Parse a single SSE line into event and data components."""
        try:
            decoded = line.decode("utf-8").strip()
            if decoded.startswith("event:"):
                return decoded.split("event:", 1)[1].strip(), None
            elif decoded.startswith("data:"):
                return None, decoded.split("data:", 1)[1].strip()
            return None, None
        except UnicodeDecodeError as e:
            logger.warning(f"Failed to decode SSE line: {e}")
            return None, None

    @classmethod
    async def iter_events(cls, response: aiohttp.ClientResponse) -> AsyncIterator[Tuple[str, str]]:
        """
        Parse SSE stream from response.

        Yields:
            Tuples of (event_type, data)
        """
        event = None
        data_buffer = []
        incomplete_line = b''

        async for chunk in response.content.iter_chunked(cls.CHUNK_SIZE):
            chunk = incomplete_line + chunk
            lines = chunk.split(b'\n')

            # Handle incomplete lines
            if not chunk.endswith(b'\n'):
                incomplete_line = lines.pop()
            else:
                incomplete_line = b''

            for line in lines:
                if not line:
                    # Empty line signals end of event
                    if event and data_buffer:
                        yield event, "\n".join(data_buffer).strip()
                    event = None
                    data_buffer = []
                    continue

                event_type, data = await cls.parse_line(line)
                if event_type:
                    event = event_type
                elif data is not None:
                    data_buffer.append(data)

        # Handle remaining data
        if incomplete_line:
            event_type, data = await cls.parse_line(incomplete_line)
            if event_type:
                event = event_type
            elif data is not None:
                data_buffer.append(data)

        if event and data_buffer:
            yield event, "\n".join(data_buffer).strip()


class ResponseProcessor:
    """Processes different types of responses from the agent."""

    @staticmethod
    def process_cortex_analyst(json_content: Dict[str, Any]) -> Dict[str, Any]:
        """Process Cortex Analyst tool results."""
        result_set = json_content.get("result_set", {})
        metadata = result_set.get("resultSetMetaData", {})
        row_types = metadata.get("rowType", [])
        columns = [row["name"] for row in row_types]
        data = result_set.get("data", [])

        source_items = []
        artifacts = []

        # Add SQL query
        sql_query = json_content.get("sql")
        if sql_query:
            source_items.append({
                "type": "GENERATED_SQL_QUERY",
                "title": "Cortex Analyst SQL Query",
                "performedQuery": sql_query
            })
            logger.info(f"Added SQL source: {sql_query[:100]}...")

        # Add records if columns exist
        if columns:
            source_items_payload = make_dates_json_safe({
                "type": "RECORDS",
                "title": "Cortex Analyst SQL Results",
                "records": {
                    "columns": columns,
                    "data": data
                }
            })

            source_items.append(source_items_payload)

            artifacts_payload = make_dates_json_safe({
                "name": "Records",
                "description": "Cortex Analyst SQL Results",
                "parts": [{
                    "type": "RECORDS",
                    "records": {
                        "columns": columns,
                        "data": data
                    }
                }]
            })
            artifacts.append(artifacts_payload)

            logger.info(
                f"Added records: {len(columns)} cols, {len(data)} rows")

        return {
            "sources": {
                "toolCallDescription": "Cortex Analyst Tool",
                "items": source_items
            },
            "artifacts": artifacts
        }

    @staticmethod
    def process_cortex_search(json_content: Dict[str, Any]) -> Dict[str, Any]:
        """Process Cortex Search tool results."""
        search_results = json_content.get("search_results", [])
        source_items = []

        for result in search_results:
            doc_title = result.get("doc_title", "Untitled")
            text_snippet = result.get("text", "")

            if text_snippet:  # Only add if there's actual content
                source_items.append({
                    "type": "SIMPLE_DOCUMENT",
                    "title": doc_title,
                    "textSnippet": text_snippet
                })

        logger.info(f"Added {len(source_items)} search results")

        return {
            "sources": {
                "toolCallDescription": "Cortex Search Tool",
                "items": source_items
            }
        }

    @staticmethod
    def process_tool_result(content: Dict[str, Any]) -> Optional[Dict[str, Any]]:
        """
        Process tool result and extract sources.

        Returns:
            Source dictionary or None if no valid data
        """
        if not isinstance(content, dict):
            logger.warning("Tool result is not a dictionary")
            return None

        tool_type = content.get("type")

        # Ensure content structure is valid
        if not (
            "content" in content
            and len(content["content"]) > 0
            and "json" in content["content"][0]
        ):
            logger.warning(
                f"Invalid content structure for tool type: {tool_type}")
            return None

        json_content = content["content"][0]["json"]

        # Process based on tool type
        if tool_type == ToolType.CORTEX_ANALYST.value and "sql" in json_content:
            return ResponseProcessor.process_cortex_analyst(json_content)

        elif tool_type == ToolType.CORTEX_SEARCH.value and "search_results" in json_content:
            return ResponseProcessor.process_cortex_search(json_content)

        else:
            logger.warning(
                f"Cannot process tool results for tool type: {tool_type}. Content: {content}")
            return None

    @staticmethod
    def format_messages(messages: List[Dict[str, str]]) -> List[Dict[str, Any]]:
        """Convert Dataiku messages to Snowflake format."""
        formatted = []
        for msg in messages:
            if msg.get("content"):
                formatted.append({
                    "role": msg["role"],
                    "content": [{"type": "text", "text": msg["content"]}]
                })
        return formatted


class SnowflakeCortexAgentLLM(BaseLLM):
    """Thread-safe Dataiku LLM implementation for Snowflake Cortex Agent."""

    def __init__(self):
        """Initialize the LLM instance."""
        self.config: Optional[AgentConfig] = None
        self.api_url: Optional[str] = None
        self.headers: Optional[Dict[str, str]] = None
        # Remove self.session - we'll create it per request instead

    def set_config(self, config: Dict[str, Any], plugin_config: Dict[str, Any]) -> None:
        """Configure the LLM with connection and model settings."""
        self.config = AgentConfig.from_dict(config)

        logger.info(
            f"Configuring agent: {self.config.agent_name} "
            f"(timeout: {self.config.api_timeout}s, retries: {self.config.max_retries})"
        )

        try:
            connector = SnowflakeConnectionHelper(
                snowflake_connection_name=self.config.connection_name)
                #snowflake_database=self.config.database,
                #snowflake_schema=self.config.schema
            #)

            self.api_url = (
                f"https://{connector.snowflake_host}/api/v2/databases/"
                f"{self.config.database}/schemas/{self.config.schema}/"
                f"agents/{self.config.agent_name}:run"
            )

            self.headers = connector.headers
            logger.info(f"Successfully configured: {self.config.agent_name}")

        except Exception as e:
            logger.error(f"Failed to configure Snowflake connection: {e}")
            raise

    async def _make_request_with_session(
        self,
        session: aiohttp.ClientSession,
        messages: List[Dict],
        retry_count: int = 0
    ) -> AsyncIterator[Dict[str, Any]]:
        """
        Make API request with a specific session.

        Args:
            session: The aiohttp ClientSession to use
            messages: Message list
            retry_count: Current retry attempt

        Yields:
            Response items
        """
        headers = self.headers.copy()
        headers.update({
            "Accept": "text/event-stream",
            "Cache-Control": "no-cache"
        })

        body = {"messages": messages}

        if self.config.debug:
            logger.debug(f"Request attempt {retry_count + 1}")

        try:
            timeout = aiohttp.ClientTimeout(
                total=self.config.api_timeout
            )

            async with session.post(
                self.api_url,
                headers=headers,
                json=body,
                timeout=timeout
            ) as resp:
                resp.raise_for_status()
                request_id = resp.headers.get("X-Snowflake-Request-Id")

                if self.config.debug:
                    logger.debug(f"Request ID: {request_id}")

                # Process SSE stream
                async for event, data_str in SSEParser.iter_events(resp):
                    result = await self._process_event(event, data_str)
                    if result["type"] != "empty":
                        yield result

        except aiohttp.ClientResponseError as e:
            if e.status >= 500 and retry_count < self.config.max_retries:
                delay = self.config.retry_delay * (retry_count + 1)
                logger.warning(
                    f"Server error {e.status}, retrying in {delay}s...")
                await asyncio.sleep(delay)
                # Recursive retry with the same session
                async for item in self._make_request_with_session(
                    session, messages, retry_count + 1
                ):
                    yield item
            else:
                logger.error(f"HTTP error: {e}")
                yield {"type": "error", "content": str(e)}

        except asyncio.TimeoutError:
            error_msg = f"Request timed out after {self.config.api_timeout} seconds"
            logger.error(error_msg)
            yield {"type": "error", "content": error_msg}

        except Exception as e:
            logger.error(f"Unexpected error: {e}")
            yield {"type": "error", "content": str(e)}

    async def _process_event(self, event: str, data_str: str) -> Dict[str, Any]:
        """Process a single SSE event."""
        if not data_str:
            return {"type": "empty", "content": None}

        try:
            data = json.loads(data_str)

            if self.config.debug:
                logger.debug(f"[{datetime.now().isoformat()}] Event: {event}")

            if event == EventType.TEXT_DELTA.value:
                text = data.get("text")
                if isinstance(text, str) and text:
                    return {"type": "text_delta", "content": text}

            elif event == EventType.THINKING_DELTA.value:
                content = data.get("text") or data.get("content")
                if isinstance(content, str) and content:
                    return {"type": "thinking", "content": content}

            elif event == EventType.TOOL_RESULT.value:
                return {"type": "tool_result", "content": data}

            elif event == EventType.RESPONSE.value:
                if self.config.debug:
                    logger.debug("Final response received")
                return {"type": "complete", "content": data}

            return {"type": "unknown", "content": data}

        except json.JSONDecodeError as e:
            logger.error(f"JSON decode error: {e}, data: {data_str[:100]}...")
            return {"type": "error", "content": f"Invalid JSON: {e}"}
        except Exception as e:
            logger.error(f"Error processing event {event}: {e}")
            return {"type": "error", "content": str(e)}

    async def aprocess_stream(
        self,
        query: SingleCompletionQuery,
        settings: CompletionSettings,
        trace: SpanBuilder
    ) -> AsyncIterator[StreamCompletionResponse]:
        """
        Thread-safe streaming response processor.
        Creates a new session for each request.

        Yields:
            Response chunks with text and metadata
        """
        # Format messages for Snowflake
        messages = ResponseProcessor.format_messages(
            [msg for msg in query["messages"] if msg.get("content")]
        )

        if self.config.debug:
            logger.debug(f"Formatted messages: {messages}")

        sources = []
        chunk_count = 0
        error_occurred = False

        # Create a new session for this specific request
        connector = aiohttp.TCPConnector(limit=100, force_close=True)
        async with aiohttp.ClientSession(
            connector=connector,
            read_bufsize=2**20  # 1MB buffer
        ) as session:
            try:
                async for response in self._make_request_with_session(session, messages):
                    chunk_count += 1

                    if self.config.debug:
                        logger.debug(
                            f"Chunk {chunk_count}: {response['type']}")

                    # Handle response types
                    if response["type"] == "text_delta":
                        chunk = DSSLLMStreamedCompletionChunk(
                            data={"text": response["content"]}
                        )
                        yield {"chunk": {"text": chunk.text}}

                    elif response["type"] == "thinking":
                        logger.info(
                            f"Thinking: {response['content'][:100]}...")

                    elif response["type"] == "tool_result":
                        sources_artifacts = ResponseProcessor.process_tool_result(
                            response["content"])
                        if sources_artifacts:
                            sources.append(sources_artifacts["sources"])
                            if "artifacts" in sources_artifacts:
                                yield {"chunk": {"type": "content", "artifacts": sources_artifacts["artifacts"]}}

                    elif response["type"] == "complete":
                        logger.info(
                            f"Stream complete after {chunk_count} chunks, {len(sources)} sources")
                        footer = DSSLLMStreamedCompletionFooter(
                            data={"additionalInformation": {
                                "sources": sources}}
                        )
                        yield footer
                        break

                    elif response["type"] == "error":
                        error_occurred = True
                        error_text = f"\n⚠️ Error: {response['content']}\n"
                        chunk = DSSLLMStreamedCompletionChunk(
                            data={"text": error_text})
                        yield {"chunk": {"text": chunk.text}}
                        break

                    # This will catch any new (and unimplemented) response types
                    else:
                        continue

                if self.config.debug:
                    logger.info(
                        f"Streaming complete. Chunks: {chunk_count}, "
                        f"Sources: {len(sources)}, Errors: {error_occurred}"
                    )

            except Exception as e:
                logger.error(f"Streaming error: {e}")
                error_text = f"\n⚠️ Error: {str(e)}\n"
                chunk = DSSLLMStreamedCompletionChunk(
                    data={"text": error_text})
                yield {"chunk": {"text": chunk.text}}

        # Session is automatically closed when exiting the async with block
        logger.debug("Session closed for this request")
