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, function_tool
import sqlite3
from datetime import datetime, timedelta
from diona.project.storage import get_storage_instance


class MilinfoLiveModel:
    """MilinfoLiveModel for processing war report messages and generating summaries."""
    
    def __init__(self, db_path: Optional[str] = None):
        """Initialize the MilinfoLiveModel.
        
        Args:
            db_path: Path to the SQLite database file. Defaults to ~/.diona/user/telegram/telegram.db
        """
        # Set default database path if not provided
        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
        
        # Get API key from environment variable
        api_key = os.environ.get("OPENAI_API_KEY")
        
        if not api_key:
            raise ValueError("OPENAI_API_KEY environment variable is required")
        
        # Use DeepSeek as the model
        self.model = "deepseek-chat"
        self.base_url = "https://api.deepseek.com/v1"
        
        # Create custom OpenAI client for DeepSeek
        self.client = AsyncOpenAI(base_url=self.base_url, api_key=api_key)
        
        # Disable tracing to avoid needing OpenAI API key for tracing
        set_tracing_disabled(disabled=True)
        
        # Create custom model provider
        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 _get_connection(self) -> sqlite3.Connection:
        """Get a connection to the SQLite database."""
        return sqlite3.connect(self.db_path)
    
    def get_recent_war_reports_tool(self):
        """Create a function tool for getting recent war reports."""
        db_path = self.db_path
        
        @function_tool
        def get_recent_war_reports(days: int = 1) -> Dict[str, Any]:
            """Get recent war report messages from the database.
            
            Args:
                days: Number of days to look back for war reports
                
            Returns:
                Dict with success status, result data, and error message
            """
            try:
                # Calculate the date threshold
                threshold_date = datetime.now() - timedelta(days=days)
                threshold_str = threshold_date.strftime("%Y-%m-%d %H:%M:%S")
                
                # Connect to the database
                conn = sqlite3.connect(db_path)
                cursor = conn.cursor()
                
                # Query for recent war reports
                query = """
                SELECT message_id, text, time 
                FROM milinfolive 
                WHERE time >= ? 
                ORDER BY time DESC
                """
                
                cursor.execute(query, (threshold_str,))
                results = cursor.fetchall()
                
                # Convert results to a list of dictionaries
                war_reports = []
                for row in results:
                    war_reports.append({
                        "message_id": row[0],
                        "text": row[1],
                        "time": row[2]
                    })
                
                conn.close()
                
                return {
                    "success": True,
                    "data": war_reports,
                    "message": f"Found {len(war_reports)} war reports from the last {days} day(s)"
                }
            
            except sqlite3.Error as e:
                return {
                    "success": False,
                    "error": f"SQL execution error: {str(e)}"
                }
            except Exception as e:
                return {
                    "success": False,
                    "error": f"Unexpected error: {str(e)}"
                }
        
        return get_recent_war_reports
    
    def get_war_report_agent(self):
        """Get an agent for retrieving war reports from the database."""
        return Agent(
            name="War Report Retriever",
            instructions="You are an agent that retrieves recent war report messages from a SQLite database. Use the get_recent_war_reports tool to fetch war reports from the last day.",
            model=self.model,
            tools=[self.get_recent_war_reports_tool()]
        )
    
    def get_analysis_agent(self):
        """Get an agent for analyzing war reports and generating summaries."""
        return Agent(
            name="War Report Analyst",
            instructions="You are an agent that analyzes war report messages and generates a comprehensive summary in a standard war report format. Focus on key events, casualties, strategic developments, and overall situation. Present the information in a clear, structured manner suitable for a daily war report bulletin.",
            model=self.model,
            tools=[]
        )
    
    async def run_war_report_analysis(self):
        """Run the complete war report analysis process.
        
        Returns:
            War report summary
        """
        # Get the first agent to retrieve war reports
        retriever_agent = self.get_war_report_agent()
        
        # Run the retriever agent to get war reports
        retrieve_result = await Runner.run(
            retriever_agent,
            "Retrieve all war report messages from the last day",
            run_config=RunConfig(model_provider=self.model_provider),
        )
        
        # Get the second agent to analyze the reports
        analyst_agent = self.get_analysis_agent()
        
        # Run the analyst agent to generate a summary
        analyze_result = await Runner.run(
            analyst_agent,
            f"Analyze the following war reports and generate a comprehensive daily war report summary in a standard format: {retrieve_result.final_output}",
            run_config=RunConfig(model_provider=self.model_provider),
        )
        
        return analyze_result.final_output
    
    def generate_war_report(self):
        """Generate a war report summary synchronously.
        
        Returns:
            War report summary
        """
        import asyncio
        return asyncio.run(self.run_war_report_analysis())