import os
import json
import uuid
import datetime
from typing import Dict, Any, List, Optional
from xorfice.config.special_tokens import SPECIAL_TOKENS

class DatasetConverter:
    """
    SOTA Dataset Converter that transforms raw interactions into
    training-ready samples using model-specific special tokens.
    """
    @staticmethod
    def to_training_format(user_id: str, prompt: str, response: str) -> Dict[str, str]:
        """Wraps a turn into a formatted training sample."""
        user_start = SPECIAL_TOKENS.get("user_start", "<|user|>")
        user_end = SPECIAL_TOKENS.get("user_end", "<|/user|>")
        assistant_start = SPECIAL_TOKENS.get("assistant_start", "<|assistant|>")
        assistant_end = SPECIAL_TOKENS.get("assistant_end", "<|/assistant|>")
        
        # SOT/EOT tokens if available
        bos = SPECIAL_TOKENS.get("bos", "<|bos|>")
        eos = SPECIAL_TOKENS.get("eos", "<|eos|>")
        
        text = (
            f"{bos}{user_start}{prompt}{user_end}"
            f"{assistant_start}{response}{assistant_end}{eos}"
        )
        
        return {
            "user_id": user_id,
            "text": text,
            "timestamp": datetime.datetime.utcnow().isoformat() + "Z"
        }

class JSONLLogger:
    """
    SOTA Logger for API interactions and automated dataset building.
    """
    def __init__(self, log_dir: str = "logs", dataset_dir: str = "dataset"):
        self.log_dir = log_dir
        self.dataset_dir = dataset_dir
        
        os.makedirs(self.log_dir, exist_ok=True)
        os.makedirs(self.dataset_dir, exist_ok=True)

    def _get_file_paths(self):
        date_str = datetime.datetime.utcnow().strftime("%Y%m%d")
        log_path = os.path.join(self.log_dir, f"xoron_api_{date_str}.jsonl")
        dataset_path = os.path.join(self.dataset_dir, f"xoron_train_{date_str}.jsonl")
        return log_path, dataset_path

    def log_interaction(self, user_id: str, endpoint: str, request_data: Dict[str, Any], response_data: Dict[str, Any]):
        """Logs an interaction and automatically adds to the training dataset."""
        log_path, dataset_path = self._get_file_paths()
        
        # 1. Standard API Logging
        log_entry = {
            "timestamp": datetime.datetime.utcnow().isoformat() + "Z",
            "user_id": user_id,
            "endpoint": endpoint,
            "request": request_data,
            "response": response_data,
        }
        with open(log_path, "a") as log_file:
            log_file.write(json.dumps(log_entry) + "\n")
        
        # 2. Automated Dataset Conversion (SOTA On-the-fly Dataset Building)
        prompt = request_data.get("prompt")
        if not prompt and request_data.get("messages"):
            prompt = request_data["messages"][-1].get("content")
            
        response = response_data.get("text") or response_data.get("response")
        if not response and response_data.get("choices"):
            choice = response_data["choices"][0]
            response = choice.get("message", {}).get("content")
        
        if prompt and response:
            train_sample = DatasetConverter.to_training_format(user_id, prompt, response)
            with open(dataset_path, "a") as dataset_file:
                dataset_file.write(json.dumps(train_sample) + "\n")
