import sys
import time
import threading
from typing import Union, List, Dict, Optional, Literal, Any
import requests
from datetime import datetime, timezone
import logging
import json
import uuid
import atexit
from pydantic import ValidationError
from threading import Timer
from raindrop.version import VERSION
from raindrop.models import TrackAIEvent, Attachment, SignalEvent, DefaultSignal, FeedbackSignal, EditSignal, PartialTrackAIEvent, PartialAIData
from raindrop.interaction import Interaction



# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

write_key = None
api_url = "https://api.raindrop.ai/v1/"
max_queue_size = 10_000
upload_size = 10
upload_interval = 1.0
buffer = []
flush_lock = threading.Lock()
debug_logs = False
flush_thread = None
shutdown_event = threading.Event()
max_ingest_size_bytes = 1 * 1024 * 1024  # 1 MB

_partial_buffers: dict[str, PartialTrackAIEvent] = {}
_partial_timers: dict[str, Timer] = {}
_PARTIAL_TIMEOUT = 2  # 2 seconds 

def set_debug_logs(value: bool):
    global debug_logs
    debug_logs = value
    if debug_logs:
        logger.setLevel(logging.DEBUG)
    else:
        logger.setLevel(logging.INFO)

def start_flush_thread():
    logger.debug("Opening flush thread")
    global flush_thread
    if flush_thread is None:
        flush_thread = threading.Thread(target=flush_loop)
        flush_thread.daemon = True
        flush_thread.start()

def flush_loop():
    while not shutdown_event.is_set():
        try:
            flush()
        except Exception as e:
            logger.error(f"Error in flush loop: {e}")
        time.sleep(upload_interval)

def flush() -> None:
    global buffer

    if buffer is None:
        logger.error("No buffer available")
        return

    logger.debug("Starting flush")

    with flush_lock:
        current_buffer = buffer
        buffer = []

    logger.debug(f"Flushing buffer size: {len(current_buffer)}")

    grouped_events = {}
    for event in current_buffer:
        endpoint = event["type"]
        data = event["data"]
        if endpoint not in grouped_events:
            grouped_events[endpoint] = []
        grouped_events[endpoint].append(data)

    for endpoint, events_data in grouped_events.items():
        for i in range(0, len(events_data), upload_size):
            batch = events_data[i:i+upload_size]
            logger.debug(f"Sending {len(batch)} events to {endpoint}")
            send_request(endpoint, batch)

    logger.debug("Flush complete")

def send_request(endpoint: str, data_entries: List[Dict[str, Union[str, Dict]]]) -> None:
    
    url = f"{api_url}{endpoint}"
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {write_key}",
    }

    max_retries = 3
    for attempt in range(max_retries):
        try:
            response = requests.post(url, json=data_entries, headers=headers)
            response.raise_for_status()
            logger.debug(f"Request successful: {response.status_code}")
            break
        except requests.exceptions.RequestException as e:
            logger.error(f"Error sending request (attempt {attempt + 1}/{max_retries}): {e}")
            if attempt == max_retries - 1:
                logger.error(f"Failed to send request after {max_retries} attempts")

def save_to_buffer(event: Dict[str, Union[str, Dict]]) -> None:
    global buffer

    if len(buffer) >= max_queue_size * 0.8:
        logger.warning(f"Buffer is at {len(buffer) / max_queue_size * 100:.2f}% capacity")

    if len(buffer) >= max_queue_size:
        logger.error("Buffer is full. Discarding event.")
        return

    logger.debug(f"Adding event to buffer: {event}")

    with flush_lock:
        buffer.append(event)

    start_flush_thread()

def identify(user_id: str, traits: Dict[str, Union[str, int, bool, float]]) -> None:
    if not _check_write_key():
        return
    data = {"user_id": user_id, "traits": traits}
    save_to_buffer({"type": "users/identify", "data": data})


def track_ai(
    user_id: str,
    event: str,
    event_id: Optional[str] = None,
    model: Optional[str] = None,
    input: Optional[str] = None,
    output: Optional[str] = None,
    convo_id: Optional[str] = None,
    properties: Optional[Dict[str, Union[str, int, bool, float]]] = None,
    timestamp: Optional[str] = None,
    attachments: Optional[List[Attachment]] = None,
) -> str:
    if not _check_write_key():
        return
    
    event_id = event_id or str(uuid.uuid4())

    try:
        payload = TrackAIEvent(
            event_id=event_id,
            user_id=user_id,
            event=event,
            timestamp=timestamp or _get_timestamp(),
            properties=properties or {},
            ai_data=dict(          # Pydantic will coerce to AIData
                model=model,
                input=input,
                output=output,
                convo_id=convo_id,
            ),
            attachments=attachments,
        )
    except ValidationError as err:
        logger.error(f"[raindrop] Invalid data passed to track_ai: {err}")
        return None


    if payload.properties is None:
        payload.properties = {}
    payload.properties["$context"] = _get_context()

    data = payload.model_dump(mode="json")

    size = _get_size(data)
    if size > max_ingest_size_bytes:
        logger.warning(
            f"[raindrop] Events larger than {max_ingest_size_bytes / (1024 * 1024)} MB may have properties truncated - "
            f"an event of size {size / (1024 * 1024):.2f} MB was logged"
        )
        return None  # Skip adding oversized events to buffer

    save_to_buffer({"type": "events/track", "data": data})
    return event_id

def shutdown():
    logger.info("Shutting down raindrop analytics")
    for eid in list(_partial_timers.keys()):
        _flush_partial_event(eid) 

    shutdown_event.set()
    if flush_thread:
        flush_thread.join(timeout=10)
    flush()  # Final flush to ensure all events are sent

def _check_write_key():
    if write_key is None:
        logger.warning("write_key is not set. Please set it before using raindrop analytics.")
        return False
    return True

def _get_context():
    return {
        "library": {
            "name": "python-sdk",
            "version": VERSION,
        },
        "metadata": {
            "pyVersion": f"v{sys.version_info.major}.{sys.version_info.minor}.{sys.version_info.micro}",
        },
    }

def _get_timestamp():
    return datetime.now(timezone.utc).replace(microsecond=0).isoformat()


def _get_size(event: dict[str, any]) -> int:
    try:
        # Add default=str to handle types like datetime
        data = json.dumps(event, default=str)
        return len(data.encode('utf-8'))
    except (TypeError, OverflowError) as e:
        logger.error(f"Error serializing event for size calculation: {e}")
        return 0 

# Signal types - This is now defined in models.py
# SignalType = Literal["default", "feedback", "edit"]

def track_signal(
    event_id: str,
    name: str,
    signal_type: Literal["default", "feedback", "edit"] = "default",
    timestamp: Optional[str] = None,
    properties: Optional[Dict[str, Any]] = None,
    attachment_id: Optional[str] = None,
    comment: Optional[str] = None,
    after: Optional[str] = None,
    sentiment: Optional[Literal["POSITIVE", "NEGATIVE"]] = None,
) -> None:
    """
    Track a signal event.

    Args:
        event_id: The ID of the event to attach the signal to
        name: Name of the signal (e.g. "thumbs_up", "thumbs_down")
        signal_type: Type of signal ("default", "feedback", or "edit")
        timestamp: Optional timestamp for the signal (ISO 8601 format)
        properties: Optional dictionary of additional properties.
        attachment_id: Optional ID of an attachment
        comment: Optional comment string (required and used only if signal_type is 'feedback').
        after: Optional after content string (required and used only if signal_type is 'edit').
        sentiment: Optional sentiment indicating if the signal is POSITIVE (default is NEGATIVE)
    """
    if not _check_write_key():
        return

    # Prepare the final properties dictionary
    final_properties = properties.copy() if properties else {}
    if signal_type == "feedback" and comment is not None:
        if "comment" in final_properties:
             logger.warning("'comment' provided as both argument and in properties; argument value used.")
        final_properties["comment"] = comment
    elif signal_type == "edit" and after is not None:
        if "after" in final_properties:
             logger.warning("'after' provided as both argument and in properties; argument value used.")
        final_properties["after"] = after

    # Prepare base arguments for all signal types
    base_args = {
        "event_id": event_id,
        "signal_name": name,
        "timestamp": timestamp or _get_timestamp(),
        "properties": final_properties,
        "attachment_id": attachment_id,
        "sentiment": sentiment
    }

    try:
        # Construct the specific signal model based on signal_type
        if signal_type == "feedback":
            payload = FeedbackSignal(**base_args, signal_type=signal_type)
        elif signal_type == "edit":
            payload = EditSignal(**base_args, signal_type=signal_type)
        else: # signal_type == "default"
            if comment is not None:
                logger.warning("'comment' argument provided for non-feedback signal type; ignored.")
            if after is not None:
                logger.warning("'after' argument provided for non-edit signal type; ignored.")
            payload = DefaultSignal(**base_args, signal_type=signal_type)

    except ValidationError as err:
        logger.error(f"[raindrop] Invalid data passed to track_signal: {err}")
        return None

    # model_dump handles the timestamp correctly
    data = payload.model_dump(mode='json')

    size = _get_size(data)
    if size > max_ingest_size_bytes:
        logger.warning(
            f"[raindrop] Events larger than {max_ingest_size_bytes / (1024 * 1024)} MB may have properties truncated - "
            f"an event of size {size / (1024 * 1024):.2f} MB was logged"
        )
        return  # Skip adding oversized events to buffer

    save_to_buffer({"type": "signals/track", "data": data})

def begin(
    user_id: str,
    event: str,
    event_id: str | None = None,
    properties: Optional[Dict[str, Any]] = None,
    input: Optional[str] = None,
    attachments: Optional[List[Attachment]] = None,
    convo_id: Optional[str] = None,
) -> Interaction:
    """
    Starts (or resumes) an interaction and returns a helper object.
    """
    eid = event_id or str(uuid.uuid4())
    
    # Prepare ai_data if input is provided
    ai_data_partial = None
    if input:
        ai_data_partial = PartialAIData(input=input, convo_id=convo_id)

    # Combine properties with initial_fields, giving precedence to initial_fields if keys clash
    final_properties = (properties or {}).copy()

    partial_event = PartialTrackAIEvent(
        event_id=eid,
        user_id=user_id,
        event=event,
        ai_data=ai_data_partial,
        properties=final_properties or None, # Pass None if empty, matching PartialTrackAIEvent defaults
        attachments=attachments,
    )

    _track_ai_partial(partial_event)
    return Interaction(eid)

def resume_interaction(event_id: str) -> Interaction:
    """Return an Interaction wrapper for an existing eventId."""
    return Interaction(event_id)

def _track_ai_partial(event: PartialTrackAIEvent) -> None:
    """
    Merge the incoming patch into an in-memory doc and flush to backend:
      • on `.finish()`  (is_pending == False)
      • or after 20 s of inactivity
    """
    eid = event.event_id

    # 1. merge
    existing = _partial_buffers.get(eid, PartialTrackAIEvent(event_id=eid))
    existing.is_pending = existing.is_pending if existing.is_pending is not None else True
    merged_dict = existing.model_dump(exclude_none=True)
    incoming = event.model_dump(exclude_none=True)

    # deep merge ai_data / properties
    def _deep(d: dict, u: dict):
        for k, v in u.items():
            d[k] = _deep(d.get(k, {}) if isinstance(v, dict) else v, v) if isinstance(v, dict) else v
        return d
    merged = _deep(merged_dict, incoming)
    merged_obj = PartialTrackAIEvent(**merged)

    _partial_buffers[eid] = merged_obj

    # 2. timer handling
    if t := _partial_timers.get(eid):
        t.cancel()
    if merged_obj.is_pending is False:
        _flush_partial_event(eid)
    else:
        _partial_timers[eid] = Timer(_PARTIAL_TIMEOUT, _flush_partial_event, args=[eid])
        _partial_timers[eid].daemon = True
        _partial_timers[eid].start()

    if debug_logs:
        logger.debug(f"[raindrop] updated partial {eid}: {merged_obj.model_dump(exclude_none=True)}")

def _flush_partial_event(event_id: str) -> None:
    """
    Send the accumulated patch as a single object to `events/track_partial`.
    """
    if t := _partial_timers.pop(event_id, None):
        t.cancel()

    evt = _partial_buffers.pop(event_id, None)
    if not evt:
        return

    # convert to ordinary TrackAIEvent-ish dict before send
    data = evt.model_dump(mode="json", exclude_none=True)
    size = _get_size(data)
    if size > max_ingest_size_bytes:
        logger.warning(f"[raindrop] partial event {event_id} > 1 MB; skipping")
        return
    
    send_request("events/track_partial", data)

atexit.register(shutdown)