from transformers import AutoModelForCausalLM, AutoTokenizer
from pydantic import BaseModel, Field
from typing import Any, Tuple, List, Dict

# Declare the thinking model
model_name = "Qwen/Qwen3-0.6B"
# load the tokenizer and the model
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype="auto")

additional_context = {
    "ONE_BR_UNITS": 23,
    "TWO_BR_UNITS": 45,
    "SOLAR_PANEL_RATING ": 625,  # W
    "BATTERY_CAPACITY": 200,  # Ah
    "BATTERY_VOLTAGE ": 96,  # V
    "SYSTEM_LOSSES ": 0.20,
    "FEED_IN_TARIFF ": 12,
    # Lighting specifications
    "LIGHTS_1BR ": 5,
    "LIGHTS_2BR ": 12,
    "LIGHT_POWER": 6,  # Watts per light
    "panel_price": 13000,
    "battery_price": 39000,
    "grid_price": 28.44,
}


def summary_generation(scenario_results: Any) -> Tuple:

    messages = [
        {
            "role": "user",
            "content": f"Analyse {scenario_results} and the {additional_context} and give thoughts on the potential investment given for the Kenyan real estate market",
        }
    ]
    text = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True,
        enable_thinking=False,  # Switches between thinking and non-thinking modes. Default is True.
    )
    model_inputs = tokenizer([text], return_tensors="pt").to(model.device)

    # conduct text completion
    generated_ids = model.generate(**model_inputs, max_new_tokens=1000)
    output_ids = generated_ids[0][len(model_inputs.input_ids[0]) :].tolist()

    # parsing thinking content
    try:
        # rindex finding 151668 (</think>)
        index = len(output_ids) - output_ids[::-1].index(151668)
    except ValueError:
        index = 0

    thinking_content = tokenizer.decode(
        output_ids[:index], skip_special_tokens=True
    ).strip("\n")
    content = tokenizer.decode(output_ids[index:], skip_special_tokens=True).strip("\n")

    loguru.info(f"thinking content:", {thinking_content})
    return ("content:", content)