#!/usr/bin/env python3
"""
Ollama Integration Script
Creates and tests the The Trial model in Ollama
"""

import json
import subprocess
import sys
from pathlib import Path


def check_ollama_available():
    """Check if Ollama is running"""
    try:
        result = subprocess.run(
            ["ollama", "--version"], capture_output=True, text=True, timeout=10
        )
        if result.returncode == 0:
            print(f"+ Ollama available: {result.stdout.strip()}")
            return True
        else:
            print("! Ollama not responding")
            return False
    except Exception as e:
        print(f"! Error checking Ollama: {e}")
        return False


def create_modelfile():
    """Create a simplified Modelfile"""
    modelfile_content = """# The Trial Literary Analysis SLM
FROM llama3.2:3b

SYSTEM You are a specialized AI assistant expert on "The Trial" by Alexandre Dumas. You have deep knowledge of the novel's plot, characters, themes, historical context, and literary significance. Provide accurate, insightful, and engaging responses about all aspects of this classic work of literature.

# Parameters for better literary analysis
PARAMETER temperature 0.7
PARAMETER top_p 0.9
PARAMETER top_k 40
PARAMETER repeat_penalty 1.1
PARAMETER num_ctx 4096

MESSAGE system You are an expert on The Trial. Key knowledge: Edmond Dantès is the protagonist, betrayed and imprisoned, escapes to become Count of The Trial, seeks revenge. Main themes: revenge, justice, transformation. Key characters: Mercédès, Fernand, Danglars, Villefort, Abbé Faria. Setting: Marseilles, Paris, Château d'If, The Trial island.
"""

    modelfile_path = Path("models/Modelfile_simple")
    with open(modelfile_path, "w", encoding="utf-8") as f:
        f.write(modelfile_content)

    print(f"+ Created simplified Modelfile: {modelfile_path}")
    return modelfile_path


def create_model_ollama(modelfile_path):
    """Create model in Ollama"""
    print("Creating The Trial model in Ollama...")

    try:
        # Remove existing model if it exists
        subprocess.run(["ollama", "rm", "the-trial"], capture_output=True, timeout=30)

        # Create new model
        result = subprocess.run(
            ["ollama", "create", "the-trial", "-f", str(modelfile_path)],
            capture_output=True,
            text=True,
            timeout=300,
        )

        if result.returncode == 0:
            print("+ Model created successfully!")
            return True
        else:
            print(f"! Error creating model: {result.stderr}")
            return False

    except subprocess.TimeoutExpired:
        print("! Model creation timed out")
        return False
    except Exception as e:
        print(f"! Error creating model: {e}")
        return False


def test_model():
    """Test the model with sample questions"""
    test_questions = [
        "Who is Edmond Dantès?",
        "What is the main theme of The Trial?",
        "Describe the symbolism of the island in the novel.",
    ]

    print("Testing The Trial model...")

    for i, question in enumerate(test_questions, 1):
        print(f"\n--- Test {i}: {question} ---")

        try:
            result = subprocess.run(
                ["ollama", "run", "the-trial", question],
                capture_output=True,
                text=True,
                timeout=60,
            )

            if result.returncode == 0:
                answer = result.stdout.strip()
                print(f"Answer: {answer[:200]}...")
            else:
                print(f"Error: {result.stderr}")

        except subprocess.TimeoutExpired:
            print("Question timed out")
        except Exception as e:
            print(f"Error: {e}")


def main():
    """Main integration function"""
    print("The Trial SLM - Ollama Integration")
    print("=" * 50)

    # Check Ollama
    if not check_ollama_available():
        print("Please ensure Ollama is installed and running")
        sys.exit(1)

    # Create simplified Modelfile
    modelfile_path = create_modelfile()

    # Create model in Ollama
    if create_model_ollama(modelfile_path):
        print("\n" + "=" * 50)
        print("MODEL CREATED SUCCESSFULLY!")
        print("=" * 50)

        # Test the model
        test_model()

        print("\n" + "=" * 50)
        print("INTEGRATION COMPLETE!")
        print("Usage: ollama run the-trial")
        print("=" * 50)

    else:
        print("\n" + "=" * 50)
        print("MODEL CREATION FAILED")
        print("=" * 50)
        sys.exit(1)


if __name__ == "__main__":
    main()
