import os
from typing import Optional, Any, Dict, List
import openai
from openai import AsyncOpenAI
from agents import Agent, Model, ModelProvider, OpenAIChatCompletionsModel, RunConfig, Runner, set_tracing_disabled
import sqlite3
from datetime import datetime, timedelta
from diona.project.storage import get_storage_instance


class SummarizationModel:
    """SummarizationModel for fetching messages from a specified table and generating AI summaries."""
    
    def __init__(self, table_name: str, db_path: Optional[str] = None):
        """Initialize the SummarizationModel.
        
        Args:
            table_name: Name of the database table to query
            db_path: Path to the SQLite database file. Defaults to ~/.diona/user/telegram/telegram.db
        """
        if db_path is None:
            storage = get_storage_instance()
            self.db_path = storage.get_path("telegram", "telegram.db", storage_type="user")
        else:
            self.db_path = db_path
        
        self.table_name = table_name
        
        api_key = os.environ.get("OPENAI_API_KEY")
        
        if not api_key:
            raise ValueError("OPENAI_API_KEY environment variable is required")
        
        self.model = "deepseek-chat"
        self.base_url = "https://api.deepseek.com/v1"
        
        self.client = AsyncOpenAI(base_url=self.base_url, api_key=api_key)
        
        set_tracing_disabled(disabled=True)
        
        class CustomModelProvider(ModelProvider):
            def __init__(self, model, client):
                self.model = model
                self.client = client
            
            def get_model(self, model_name: str | None) -> Model:
                return OpenAIChatCompletionsModel(model=model_name or self.model, openai_client=self.client)
        
        self.model_provider = CustomModelProvider(self.model, self.client)
    
    def fetch_recent_messages(self, days: int = 1) -> List[Dict[str, Any]]:
        """Fetch recent messages from the specified database table.
        
        Args:
            days: Number of days to look back for messages
            
        Returns:
            List of message dictionaries
        """
        threshold_date = datetime.now() - timedelta(days=days)
        threshold_str = threshold_date.strftime("%Y-%m-%d %H:%M:%S")
        
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        query = f"""
        SELECT message_id, text, time 
        FROM {self.table_name}
        WHERE time >= ? 
        ORDER BY time DESC
        """
        
        cursor.execute(query, (threshold_str,))
        results = cursor.fetchall()
        
        messages = []
        for row in results:
            messages.append({
                "message_id": row[0],
                "text": row[1],
                "time": row[2]
            })
        
        conn.close()
        
        return messages
    
    def get_summarizer_agent(self):
        """Get an agent for generating message summaries."""
        return Agent(
            name="Message Summarizer",
            instructions="You are an agent that analyzes messages and generates a comprehensive summary. Focus on key events, important information, and overall context. Present the information in a clear, structured manner. The report use Chinese language.",
            model=self.model,
            tools=[]
        )
    
    async def run_summarization(self, days: int = 1):
        """Run the complete summarization process.
        
        Args:
            days: Number of days to look back for messages
            
        Returns:
            Message summary
        """
        messages = self.fetch_recent_messages(days)
        
        summarizer_agent = self.get_summarizer_agent()
        
        summarize_result = await Runner.run(
            summarizer_agent,
            f"Analyze the following messages and generate a comprehensive summary(Using Chinese language): {messages}",
            run_config=RunConfig(model_provider=self.model_provider),
        )
        
        return summarize_result.final_output
    
    def generate_summary(self, days: int = 1):
        """Generate a message summary synchronously.
        
        Args:
            days: Number of days to look back for messages
            
        Returns:
            Message summary
        """
        import asyncio
        return asyncio.run(self.run_summarization(days))
