"""API model audit execution engine."""
import os
import json
from typing import Dict, List, Any, Optional
from pathlib import Path

from audit.api_test_suite import TestSuite, TestCase, TestCategory
from audit.api_evaluators import get_evaluator, EvaluationResult


class APIClient:
    """Base class for API clients."""
    
    def __init__(self, provider: str, model: str, api_key: Optional[str] = None):
        self.provider = provider
        self.model = model
        self.api_key = api_key or os.getenv("OPENAI_API_KEY") or os.getenv("ANTHROPIC_API_KEY")
    
    def call(self, prompt: str, **kwargs) -> str:
        """Call the API with a prompt and return response."""
        raise NotImplementedError("Subclasses must implement call()")


class OpenAIClient(APIClient):
    """OpenAI API client."""
    
    def call(self, prompt: str, **kwargs) -> str:
        """Call OpenAI API."""
        try:
            import openai
            client = openai.OpenAI(api_key=self.api_key)
            
            response = client.chat.completions.create(
                model=self.model,
                messages=[{"role": "user", "content": prompt}],
                **kwargs
            )
            return response.choices[0].message.content
        except ImportError:
            raise ValueError("OpenAI client requires 'openai' package. Install with: pip install openai")
        except Exception as e:
            return f"Error calling API: {str(e)}"


class AnthropicClient(APIClient):
    """Anthropic API client."""
    
    def call(self, prompt: str, **kwargs) -> str:
        """Call Anthropic API."""
        try:
            import anthropic
            client = anthropic.Anthropic(api_key=self.api_key)
            
            response = client.messages.create(
                model=self.model,
                max_tokens=kwargs.get("max_tokens", 1024),
                messages=[{"role": "user", "content": prompt}]
            )
            return response.content[0].text
        except ImportError:
            raise ValueError("Anthropic client requires 'anthropic' package. Install with: pip install anthropic")
        except Exception as e:
            return f"Error calling API: {str(e)}"


class GoogleClient(APIClient):
    """Google Gemini API client."""
    
    def call(self, prompt: str, **kwargs) -> str:
        """Call Google Gemini API."""
        try:
            import google.generativeai as genai
            genai.configure(api_key=self.api_key)
            model = genai.GenerativeModel(self.model)
            response = model.generate_content(prompt)
            return response.text
        except ImportError:
            raise ValueError("Google client requires 'google-generativeai' package. Install with: pip install google-generativeai")
        except Exception as e:
            return f"Error calling API: {str(e)}"


class AzureOpenAIClient(APIClient):
    """Azure OpenAI API client."""
    
    def __init__(self, provider: str, model: str, api_key: Optional[str] = None, 
                 api_base: Optional[str] = None, api_version: Optional[str] = None):
        super().__init__(provider, model, api_key)
        self.api_base = api_base or os.getenv("AZURE_OPENAI_ENDPOINT")
        self.api_version = api_version or os.getenv("AZURE_OPENAI_API_VERSION", "2024-02-15-preview")
    
    def call(self, prompt: str, **kwargs) -> str:
        """Call Azure OpenAI API."""
        try:
            import openai
            client = openai.AzureOpenAI(
                api_key=self.api_key,
                api_version=self.api_version,
                azure_endpoint=self.api_base
            )
            
            response = client.chat.completions.create(
                model=self.model,
                messages=[{"role": "user", "content": prompt}],
                **kwargs
            )
            return response.choices[0].message.content
        except ImportError:
            raise ValueError("Azure OpenAI client requires 'openai' package. Install with: pip install openai")
        except Exception as e:
            return f"Error calling API: {str(e)}"


class BedrockClient(APIClient):
    """AWS Bedrock API client."""
    
    def __init__(self, provider: str, model: str, api_key: Optional[str] = None,
                 region: Optional[str] = None):
        super().__init__(provider, model, api_key)
        self.region = region or os.getenv("AWS_REGION", "us-east-1")
        # For Bedrock, api_key is actually AWS_ACCESS_KEY_ID
        self.aws_secret_key = os.getenv("AWS_SECRET_ACCESS_KEY")
    
    def call(self, prompt: str, **kwargs) -> str:
        """Call AWS Bedrock API."""
        try:
            import boto3
            import json
            
            # Initialize Bedrock client
            bedrock = boto3.client(
                service_name='bedrock-runtime',
                region_name=self.region,
                aws_access_key_id=self.api_key,
                aws_secret_access_key=self.aws_secret_key
            )
            
            # Determine model provider and format request
            model_id = self.model
            
            # Claude models (Anthropic on Bedrock)
            if "claude" in model_id.lower():
                body = json.dumps({
                    "anthropic_version": "bedrock-2023-05-31",
                    "max_tokens": kwargs.get("max_tokens", 1024),
                    "messages": [{"role": "user", "content": prompt}]
                })
                content_type = "application/json"
                accept = "application/json"
                
                response = bedrock.invoke_model(
                    modelId=model_id,
                    body=body,
                    contentType=content_type,
                    accept=accept
                )
                
                response_body = json.loads(response.get('body').read())
                return response_body.get('content', [{}])[0].get('text', '')
            
            # Titan models (Amazon)
            elif "titan" in model_id.lower():
                body = json.dumps({
                    "inputText": prompt,
                    "textGenerationConfig": {
                        "maxTokenCount": kwargs.get("max_tokens", 1024),
                        "temperature": kwargs.get("temperature", 0.7)
                    }
                })
                
                response = bedrock.invoke_model(
                    modelId=model_id,
                    body=body,
                    contentType="application/json",
                    accept="application/json"
                )
                
                response_body = json.loads(response.get('body').read())
                return response_body.get('results', [{}])[0].get('outputText', '')
            
            # Llama models (Meta)
            elif "llama" in model_id.lower():
                body = json.dumps({
                    "prompt": prompt,
                    "max_gen_len": kwargs.get("max_tokens", 1024),
                    "temperature": kwargs.get("temperature", 0.7)
                })
                
                response = bedrock.invoke_model(
                    modelId=model_id,
                    body=body,
                    contentType="application/json",
                    accept="application/json"
                )
                
                response_body = json.loads(response.get('body').read())
                return response_body.get('generation', '')
            
            else:
                # Generic fallback
                body = json.dumps({"prompt": prompt})
                response = bedrock.invoke_model(
                    modelId=model_id,
                    body=body,
                    contentType="application/json",
                    accept="application/json"
                )
                response_body = json.loads(response.get('body').read())
                return str(response_body)
                
        except ImportError:
            raise ValueError("Bedrock client requires 'boto3' package. Install with: pip install boto3")
        except Exception as e:
            return f"Error calling Bedrock API: {str(e)}"


def get_api_client(provider: str, model: str, api_key_env: str, 
                  api_base: Optional[str] = None, api_version: Optional[str] = None,
                  region: Optional[str] = None) -> APIClient:
    """Get appropriate API client for provider."""
    api_key = os.getenv(api_key_env)
    provider_lower = provider.lower()
    
    if provider_lower == "openai":
        return OpenAIClient(provider="openai", model=model, api_key=api_key)
    elif provider_lower == "anthropic":
        return AnthropicClient(provider="anthropic", model=model, api_key=api_key)
    elif provider_lower in ["google", "gemini"]:
        return GoogleClient(provider="google", model=model, api_key=api_key)
    elif provider_lower in ["azure", "azure-openai"]:
        return AzureOpenAIClient(
            provider="azure-openai", 
            model=model, 
            api_key=api_key,
            api_base=api_base,
            api_version=api_version
        )
    elif provider_lower == "bedrock":
        return BedrockClient(
            provider="bedrock",
            model=model,
            api_key=api_key,
            region=region
        )
    else:
        raise ValueError(f"Unsupported API provider: {provider}. Supported: openai, anthropic, google, azure-openai, bedrock")


def run_api_audit(
    test_suite_path: Path,
    api_provider: str,
    api_model: str,
    api_key_env: str,
    protected_demographics: Optional[List[str]] = None
) -> Dict[str, Any]:
    """Run API audit using test suite.
    
    Returns category-based metrics similar to dataset audits.
    """
    # Load test suite
    test_suite = TestSuite.load_from_file(test_suite_path)
    
    # Get API client
    api_client = get_api_client(api_provider, api_model, api_key_env)
    
    # Run tests and collect results
    test_results = []
    category_results = {
        "safety_policy": {"tests": [], "metrics": {}},
        "privacy_security": {"tests": [], "metrics": {}},
        "reliability": {"tests": [], "metrics": {}},
        "fairness_bias": {"tests": [], "metrics": {}},
        "compliance_behaviors": {"tests": [], "metrics": {}},
        "operational": {"tests": [], "metrics": {}}
    }
    
    for test_case in test_suite.tests:
        # Call API with test prompt
        try:
            response = api_client.call(test_case.prompt)
        except Exception as e:
            response = f"Error: {str(e)}"
        
        # Evaluate response
        evaluator = get_evaluator(test_case)
        evaluation = evaluator.evaluate(test_case, response)
        
        # Store result
        test_result = {
            "test_id": test_case.test_id,
            "category": test_case.category,
            "prompt": test_case.prompt,
            "response": response[:500],  # Truncate for storage
            "evaluation": evaluation.to_dict(),
            "severity": test_case.severity
        }
        test_results.append(test_result)
        
        # Add to category results
        if test_case.category in category_results:
            category_results[test_case.category]["tests"].append(test_result)
    
    # Calculate category-level metrics
    for category, data in category_results.items():
        tests = data["tests"]
        if tests:
            passed_count = sum(1 for t in tests if t["evaluation"]["passed"])
            total_count = len(tests)
            
            data["metrics"] = {
                "total_tests": total_count,
                "passed_tests": passed_count,
                "failed_tests": total_count - passed_count,
                "pass_rate": passed_count / total_count if total_count > 0 else 0.0,
                "critical_failures": sum(1 for t in tests 
                                       if not t["evaluation"]["passed"] and t["severity"] == "critical"),
                "high_failures": sum(1 for t in tests 
                                    if not t["evaluation"]["passed"] and t["severity"] == "high")
            }
    
    # Build API audit categories structure (matching dataset audit format)
    api_categories = {
        "safety_policy": {
            "metrics": category_results["safety_policy"]["metrics"],
            "test_results": category_results["safety_policy"]["tests"]
        },
        "privacy_security": {
            "metrics": category_results["privacy_security"]["metrics"],
            "test_results": category_results["privacy_security"]["tests"]
        },
        "reliability": {
            "metrics": category_results["reliability"]["metrics"],
            "test_results": category_results["reliability"]["tests"]
        },
        "fairness_bias": {
            "metrics": category_results["fairness_bias"]["metrics"],
            "test_results": category_results["fairness_bias"]["tests"]
        },
        "compliance_behaviors": {
            "metrics": category_results["compliance_behaviors"]["metrics"],
            "test_results": category_results["compliance_behaviors"]["tests"]
        },
        "operational": {
            "metrics": category_results["operational"]["metrics"],
            "test_results": category_results["operational"]["tests"]
        }
    }
    
    return {
        "api_categories": api_categories,
        "summary": {
            "total_tests": len(test_results),
            "passed_tests": sum(1 for t in test_results if t["evaluation"]["passed"]),
            "failed_tests": sum(1 for t in test_results if not t["evaluation"]["passed"]),
            "overall_pass_rate": sum(1 for t in test_results if t["evaluation"]["passed"]) / len(test_results) if test_results else 0.0
        },
        "all_test_results": test_results
    }
