import os
import boto3
import json
from dotenv import load_dotenv

load_dotenv()

class NovaProAgent:
    def __init__(self):
        print("NovaProAgent initialized.")
        
        # Get AWS credentials from environment variables
        aws_access_key_id = os.getenv('AWS_ACCESS_KEY_ID')
        aws_secret_access_key = os.getenv('AWS_SECRET_ACCESS_KEY')

        # Initialize the AWS client
        boto3.client(
            's3',
            aws_access_key_id=aws_access_key_id,
            aws_secret_access_key=aws_secret_access_key
        )
        session = boto3.session.Session()
        
        self.bedrock_client = boto3.client(
            service_name='bedrock-runtime',
            region_name='us-east-1'
        )

        self.model_id = "amazon.nova-pro-v1:0"
        self.content_type = "application/json"
        self.accept = "application/json"
        
    async def __call__(self, question: str) -> str:
        print(f"NovaProAgent received question (first 50 chars): {question[:50]}...")
        
        try:
            # Detect question type and adjust approach
            needs_reasoning = any(keyword in question.lower() for keyword in [
                'calculate', 'how many', 'what is the', 'find', 'determine', 'solve',
                'table', 'given', 'prove', 'counter-example'
            ])
            
            if needs_reasoning:
                prompt = f"""You are an expert problem solver. Think step by step to solve this question accurately.

Question: {question}

Think through this step by step, then provide your final answer:"""
                max_tokens = 300
            else:
                prompt = f"""Answer this question directly and concisely. Provide only the essential information requested.

Question: {question}

Answer:"""
                max_tokens = 150
            
            # Prepare the request payload for Nova Pro
            payload = {
                "messages": [
                    {
                        "role": "user",
                        "content": [{
                            "text": prompt
                        }]
                    }
                ],
                "inferenceConfig": {
                    "max_new_tokens": max_tokens,
                    "temperature": 0.1 if needs_reasoning else 0.0
                }
            }
            
            # Call Nova Lite model
            response = self.bedrock_client.invoke_model(
                modelId=self.model_id,
                contentType=self.content_type,
                accept=self.accept,
                body=json.dumps(payload)
            )
            
            # Parse response
            response_body = json.loads(response['body'].read())
            answer = response_body['output']['message']['content'][0]['text']
            
            # Clean up the answer
            answer = answer.strip()
            
            # Handle questions requiring external resources
            if any(phrase in question.lower() for phrase in [
                'attached', 'video', 'image', 'audio', 'file', 'excel', '.mp3', '.jpg', '.png'
            ]):
                if 'video' in question.lower() or 'audio' in question.lower():
                    return "I cannot access external media files."
                elif 'image' in question.lower():
                    return "I cannot view images."
                elif 'excel' in question.lower() or 'file' in question.lower():
                    return "I cannot access attached files."
            
            # Extract final answer if reasoning was used
            if needs_reasoning and 'final answer' in answer.lower():
                lines = answer.split('\n')
                for line in reversed(lines):
                    if 'final answer' in line.lower() or 'answer:' in line.lower():
                        # Extract the part after the colon
                        if ':' in line:
                            answer = line.split(':', 1)[1].strip()
                        break
            
            # Remove verbose beginnings
            verbose_starts = [
                "To answer this question", "Based on the information", "According to",
                "The answer is", "Looking at", "Step by step", "Let me think"
            ]
            
            for start in verbose_starts:
                if answer.lower().startswith(start.lower()):
                    sentences = answer.split('. ')
                    for sentence in sentences[1:]:
                        if len(sentence.strip()) > 5:
                            answer = sentence.strip()
                            break
            
            # Clean up common patterns
            answer = answer.replace('**', '').replace('*', '')
            
            return answer
            
        except Exception as e:
            print(f"Error calling Nova Lite: {e}")
            return "Unable to process request."