import anthropic
import os
from dotenv import load_dotenv
import json
from src.prompt import advanced_summary_prompt, basic_summary_prompt, custom_template_prompt
load_dotenv()

def general_summary(transcription, plan_tier="pro"):
    """
    Generate a summary of the transcription based on the user's plan tier.

    Args:
        transcription: The transcription to summarize
        plan_tier: The user's plan tier ("freemium" or "pro")

    Returns:
        A JSON object containing the summary
    """
    client = anthropic.Anthropic(
        api_key=os.getenv("ANTHTROPIC_API_KEY"),
    )

    # Select the appropriate prompt based on the user's plan tier
    if plan_tier.lower() == "freemium":
        prompt = basic_summary_prompt
        max_tokens = 2000  # Reduced token count for basic summaries
    else:  # Default to pro
        prompt = advanced_summary_prompt
        max_tokens = 4000

    message = client.messages.create(
        model="claude-3-5-sonnet-20241022",
        max_tokens=max_tokens,
        messages=[
            {"role": "user", "content": f"{prompt}"},
            {"role": "user", "content": f"Transcription: {transcription}"}
        ]
    )

    text = message.content[0].text
    return json.loads(text)




def custom_summary(template, transcription):
    client = anthropic.Anthropic(
        api_key=os.getenv("ANTHTROPIC_API_KEY"),
    )
    message = client.messages.create(
        model="claude-3-5-sonnet-20241022",
        max_tokens=8000,
        messages=[
            {"role": "user", "content": f"{custom_template_prompt}"},
            {"role": "user", "content": f"TEMPLATE : {template}"},
            {"role": "user", "content": f"Transcription: {transcription}"}
        ]
    )

    text = message.content[0].text
    return json.loads(text)