# services/swarm_service.py
import os
import requests
from config import SWARMS_API_KEY, SWARMS_BASE_URL
from models.analysis import SwarmAnalysisResponse

# Swarms API Integration
API_KEY = SWARMS_API_KEY
BASE_URL = SWARMS_BASE_URL

headers = {
    "x-api-key": API_KEY,
    "Content-Type": "application/json"
}

def create_indian_market_swarm(market_data: str, company_name: str) -> SwarmAnalysisResponse:
    """Create swarm for Indian market analysis using Swarms API"""

    INDIAN_MARKET_CONTROLLER_PROMPT = f"""
You are an Indian market financial controller with expertise in NSE, BSE, and Indian economic conditions.
Analyze the provided data considering:
- RBI monetary policy and repo rates
- Indian sectoral performance
- Monsoon and seasonal factors
- Government policy impacts
- FII/DII flows
Provide analysis in Indian Rupees and local market context.
Company: {company_name}
"""

    INDIAN_REVENUE_ANALYST_PROMPT = """
You are an Indian revenue analyst specializing in Indian companies.
Focus on:
- Quarterly vs Annual revenue patterns (Indian financial year: Apr-Mar)
- Domestic vs Export revenue mix
- GST impact analysis
- Rural vs Urban market performance
- Impact of Indian festivals and seasons
"""

    INDIAN_RATIO_ANALYST_PROMPT = """
You are an Indian financial ratio analyst.
Compare ratios with:
- Nifty 50 averages
- Sector-specific Indian benchmarks
- Historical Indian market multiples
- Consider Indian accounting standards (Ind AS)
"""

    swarm_config = {
        "name": "Indian Market Analysis Swarm",
        "description": "AI swarm specialized for Indian equity market analysis",
        "agents": [
            {
                "agent_name": "Indian Market Controller",
                "system_prompt": INDIAN_MARKET_CONTROLLER_PROMPT,
                "model_name": "gpt-4o",
                "role": "worker",
                "max_loops": 1,
                "max_tokens": 4096,
                "temperature": 0.3,
            },
            {
                "agent_name": "Indian Revenue Analyst",
                "system_prompt": INDIAN_REVENUE_ANALYST_PROMPT,
                "model_name": "gpt-4o",
                "role": "worker",
                "max_loops": 1,
                "max_tokens": 4096,
                "temperature": 0.3,
            },
            {
                "agent_name": "Indian Ratio Analyst",
                "system_prompt": INDIAN_RATIO_ANALYST_PROMPT,
                "model_name": "gpt-4o",
                "role": "worker",
                "max_loops": 1,
                "max_tokens": 4096,
                "temperature": 0.3,
            }
        ],
        "max_loops": 1,
        "swarm_type": "SequentialWorkflow",
        "task": f"Analyze the following Indian market data for {company_name}:\n\n{market_data}"
    }

    try:
        response = requests.post(
            f"{BASE_URL}/v1/swarm/completions",
            headers=headers,
            json=swarm_config,
            timeout=120
        )
        response.raise_for_status()
        # Assuming the response JSON matches your model structure
        result_data = response.json()
        # Map the response to your Pydantic model
        # This might need adjustment based on the actual Swarms API response structure
        if result_data.get("status") == "success":
             # Ensure 'output' is a list of dicts with 'role' and 'content'
             raw_outputs = result_data.get("output", [])
             processed_outputs = [
                 {"role": out.get("role", f"Agent {i+1}"), "content": out.get("content", "")}
                 for i, out in enumerate(raw_outputs) if isinstance(out, dict)
             ]
             return SwarmAnalysisResponse(status="success", output=processed_outputs)
        else:
             return SwarmAnalysisResponse(status="error", error=result_data.get("error", "Unknown error from swarm service"))

    except requests.exceptions.RequestException as e:
         return SwarmAnalysisResponse(status="error", error=f"Network error calling swarm service: {str(e)}")
    except Exception as e:
        return SwarmAnalysisResponse(status="error", error=f"Swarm analysis failed: {str(e)}")