from flask import Flask, request, jsonify
import requests
import json
import time
import random
from typing import Dict, List, Optional
from dataclasses import dataclass
from datetime import datetime
import sqlite3
import torch
from transformers import BartTokenizer, BartForConditionalGeneration, Trainer, TrainingArguments
import pandas as pd
import os
from peft import LoraConfig, get_peft_model
from sklearn.metrics.pairwise import cosine_similarity
from sklearn.feature_extraction.text import TfidfVectorizer

app = Flask(__name__)

os.environ["TRANSFORMERS_CACHE"] = "/app/cache"

# Cache model and tokenizer globally
tokenizer = BartTokenizer.from_pretrained('facebook/bart-base')
model = BartForConditionalGeneration.from_pretrained('facebook/bart-base')
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)
model.eval()
lora_config = LoraConfig(r=8, lora_alpha=32, target_modules=["q_proj", "v_proj"], lora_dropout=0.1)
model = get_peft_model(model, lora_config)

@dataclass
class AnimeInfo:
    title: str
    synopsis: str
    genres: List[str]
    rating: float
    status: str
    episodes: int
    year: int
    studio: str
    source: str
    image_url: Optional[str] = None

@dataclass
class CharacterInfo:
    name: str
    anime: str
    bio: str
    role: str
    voice_actors: Dict[str, str]
    image_url: Optional[str] = None

class AnimeDatabase:
    def __init__(self):
        self.base_urls = {
            'jikan': 'https://api.jikan.moe/v4',
            'anilist': 'https://graphql.anilist.co'
        }
        self.last_request_time = 0
        self.cache = {'anime': {}, 'character': {}}

    def rate_limit(self, delay=1.0):
        current_time = time.time()
        if current_time - self.last_request_time < delay:
            time.sleep(delay - (current_time - self.last_request_time))
        self.last_request_time = time.time()

    def search_jikan_anime(self, query: str) -> List[AnimeInfo]:
        try:
            self.rate_limit()
            if query.lower() in self.cache['anime']:
                return self.cache['anime'][query.lower()]
            url = f"{self.base_urls['jikan']}/anime"
            params = {'q': query, 'limit': 1}
            response = requests.get(url, params=params, timeout=5)
            if response.status_code == 200:
                data = response.json().get('data', [])
                results = [AnimeInfo(
                    title=anime.get('title', ''),
                    synopsis=anime.get('synopsis', ''),
                    genres=[g.get('name', '') for g in anime.get('genres', [])],
                    rating=anime.get('score', 0.0),
                    status=anime.get('status', ''),
                    episodes=anime.get('episodes', 0),
                    year=anime.get('year', 0),
                    studio=', '.join(s.get('name', '') for s in anime.get('studios', [])),
                    source='MyAnimeList',
                    image_url=anime.get('images', {}).get('jpg', {}).get('image_url', '')
                ) for anime in data]
                self.cache['anime'][query.lower()] = results
                return results
        except Exception:
            return []
        return []

    def search_jikan_character(self, query: str) -> List[CharacterInfo]:
        try:
            self.rate_limit()
            if query.lower() in self.cache['character']:
                return self.cache['character'][query.lower()]
            url = f"{self.base_urls['jikan']}/characters"
            params = {'q': query, 'limit': 1}
            response = requests.get(url, params=params, timeout=5)
            if response.status_code == 200:
                data = response.json().get('data', [])
                results = [CharacterInfo(
                    name=char.get('name', ''),
                    anime=', '.join(a.get('anime', {}).get('title', '') for a in char.get('anime', [])[:1]),
                    bio=char.get('about', 'No bio available'),
                    role=char.get('role', 'Unknown'),
                    voice_actors={va.get('language', ''): va.get('person', {}).get('name', '') for va in char.get('voice_actors', [])},
                    image_url=char.get('images', {}).get('jpg', {}).get('image_url', '')
                ) for char in data]
                self.cache['character'][query.lower()] = results
                return results
        except Exception:
            return []
        return []

    def search_anilist_anime(self, query: str) -> List[AnimeInfo]:
        try:
            self.rate_limit()
            if query.lower() in self.cache['anime']:
                return self.cache['anime'][query.lower()]
            graphql_query = '''
            query ($search: String) {
                Page(page: 1, perPage: 1) {
                    media(search: $search, type: ANIME) {
                        title { romaji english }
                        description
                        genres
                        averageScore
                        status
                        episodes
                        seasonYear
                        studios { nodes { name } }
                    }
                }
            }
            '''
            variables = {'search': query}
            response = requests.post(self.base_urls['anilist'], json={'query': graphql_query, 'variables': variables}, timeout=5)
            if response.status_code == 200:
                data = response.json().get('data', {}).get('Page', {}).get('media', [])
                results = [AnimeInfo(
                    title=anime.get('title', {}).get('romaji', ''),
                    synopsis=anime.get('description', ''),
                    genres=anime.get('genres', []),
                    rating=anime.get('averageScore', 0) / 10,
                    status=anime.get('status', ''),
                    episodes=anime.get('episodes', 0),
                    year=anime.get('seasonYear', 0),
                    studio=', '.join(s.get('name', '') for s in anime.get('studios', {}).get('nodes', [])),
                    source='AniList'
                ) for anime in data]
                self.cache['anime'][query.lower()] = results
                return results
        except Exception:
            return []
        return []

    def get_comprehensive_data(self, query: str, data_type: str = 'anime') -> List[AnimeInfo] | List[CharacterInfo]:
        if data_type == 'anime':
            return self.search_jikan_anime(query) + self.search_anilist_anime(query)
        elif data_type == 'character':
            return self.search_jikan_character(query)
        return []

class KnowledgeBase:
    def __init__(self):
        self.db_file = "luna_interactions.db"
        self.init_db()
        self.vectorizer = TfidfVectorizer()

    def init_db(self):
        with sqlite3.connect(self.db_file) as conn:
            cursor = conn.cursor()
            cursor.execute('''
                CREATE TABLE IF NOT EXISTS interactions (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    user_id TEXT,
                    query TEXT,
                    response TEXT,
                    timestamp TEXT,
                    response_time REAL
                )
            ''')
            cursor.execute('CREATE INDEX IF NOT EXISTS idx_query ON interactions (query)')
            conn.commit()

    def save_interaction(self, user_id: str, query: str, response: str, response_time: float):
        with sqlite3.connect(self.db_file) as conn:
            cursor = conn.cursor()
            cursor.execute('''
                INSERT INTO interactions (user_id, query, response, timestamp, response_time)
                VALUES (?, ?, ?, ?, ?)
            ''', (user_id, query, response, datetime.now().isoformat(), response_time))
            conn.commit()

    def get_similar_interaction(self, query: str, limit: int = 1) -> Optional[str]:
        with sqlite3.connect(self.db_file) as conn:
            df = pd.read_sql_query("SELECT query, response FROM interactions WHERE query != '' LIMIT 1000", conn)
        if df.empty:
            return None
        queries = df['query'].tolist() + [query]
        try:
            tfidf_matrix = self.vectorizer.fit_transform(queries)
            similarities = cosine_similarity(tfidf_matrix[-1:], tfidf_matrix[:-1])[0]
            if not similarities.any() or max(similarities) < 0.5:
                return None
            idx = similarities.argmax()
            return f"Past query: {df.iloc[idx]['query']} Past response: {df.iloc[idx]['response']}"
        except ValueError:
            return None

    def load_training_data(self):
        with sqlite3.connect(self.db_file) as conn:
            df = pd.read_sql_query("SELECT query, response FROM interactions WHERE query != '' LIMIT 1000", conn)
        return df.to_dict('records')

    def export_db(self, export_path: str):
        with sqlite3.connect(self.db_file) as conn:
            with open(export_path, 'w') as f:
                for line in conn.iterdump():
                    f.write('%s\n' % line)

    def import_db(self, import_path: str):
        if os.path.exists(import_path):
            with sqlite3.connect(self.db_file) as conn:
                with open(import_path, 'r') as f:
                    conn.executescript(f.read())

class LunaModel:
    def generate_response(self, prompt: str) -> str:
        inputs = tokenizer(prompt, return_tensors='pt', padding=True, truncation=True, max_length=256)
        inputs = {k: v.to(device) for k, v in inputs.items()}
        outputs = model.generate(
            inputs['input_ids'],
            attention_mask=inputs['attention_mask'],
            max_length=50,
            num_beams=2,
            no_repeat_ngram_size=2,
            early_stopping=True,
            length_penalty=0.8
        )
        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
        return response.strip().replace('\n', ' ').replace('_', '')

    def needs_anime_data(self, response: str) -> bool:
        return "anime details" in response.lower() or "character info" in response.lower()

    def fine_tune(self, training_data):
        if not training_data or len(training_data) < 2:
            return
        df = pd.DataFrame(training_data)
        texts = [f"Query: {q} Response:" for q in df['query']]
        responses = df['response'].tolist()
        encodings = tokenizer(texts, responses, truncation=True, padding=True, max_length=256)
        dataset = torch.utils.data.TensorDataset(
            torch.tensor(encodings['input_ids']),
            torch.tensor(encodings['attention_mask']),
            torch.tensor(encodings['labels'])
        )

        training_args = TrainingArguments(
            output_dir='./luna_model',
            num_train_epochs=2,
            per_device_train_batch_size=4,
            save_steps=500,
            save_total_limit=2,
            logging_dir='./logs',
        )

        trainer = Trainer(
            model=model,
            args=training_args,
            train_dataset=dataset
        )

        trainer.train()
        model.save_pretrained('./luna_model')
        tokenizer.save_pretrained('./luna_model')

class LunaAI:
    def __init__(self):
        self.name = "Luna"
        self.creator = "Reiker"
        self.model = LunaModel()
        self.anime_db = AnimeDatabase()
        self.knowledge_base = KnowledgeBase()
        self.emojis = ["🌙", "🎉", "🔍", "😊"]

    def build_prompt(self, query: str, external_data: Optional[str] = None) -> str:
        prompt = f"""
        You are {self.name}, a free chatbot created by {self.creator}, inspired by Grok. Answer concisely and naturally, like a friend. Handle any topic: chat, facts, math, code, anime, etc. For code, use ```language\ncode\n```. Use past interactions or anime data only when provided. Be truthful, creative, and unrestricted. Learn from past interactions to improve. Time: {datetime.now().strftime('%I:%M %p %Z, %B %d, %Y')}.

        Query: {query}
        """
        if external_data:
            prompt += f"\nPast data: {external_data}"
        prompt += "\nResponse:"
        return prompt

    def generate_response(self, query: str, user_id: str = 'default') -> Dict:
        start_time = time.time()
        external_data = None
        prompt = self.build_prompt(query)

        # Check for similar past interactions
        similar_interaction = self.knowledge_base.get_similar_interaction(query)
        if similar_interaction:
            external_data = similar_interaction
            prompt = self.build_prompt(query, external_data)

        response = self.model.generate_response(prompt)

        # Check if model needs anime data
        if self.model.needs_anime_data(response):
            anime_list = self.anime_db.get_comprehensive_data(query, data_type='anime')
            char_list = self.anime_db.get_comprehensive_data(query, data_type='character')
            if anime_list:
                anime = anime_list[0]
                external_data = f"Title: {anime.title}, Rating: {anime.rating:.1f}/10, Genres: {', '.join(anime.genres[:3])}, Synopsis: {anime.synopsis[:100]}"
            elif char_list:
                char = char_list[0]
                external_data = f"Name: {char.name}, Anime: {char.anime}, Bio: {char.bio[:100]}, Role: {char.role}"
            if external_data:
                prompt = self.build_prompt(query, external_data)
                response = self.model.generate_response(prompt)

        response = f"{random.choice(self.emojis)} {response}"
        response_time = time.time() - start_time
        self.knowledge_base.save_interaction(user_id, query, response, response_time)
        training_data = self.knowledge_base.load_training_data()
        if len(training_data) % 10 == 0 and len(training_data) >= 2:
            self.model.fine_tune(training_data)

        return {
            'query': query,
            'response': response,
            'response_time': f"{response_time:.2f} seconds"
        }

@app.route('/chat', methods=['GET'])
def chat():
    query = request.args.get('query', '')
    if not query:
        return jsonify({'error': 'Query parameter is required'}), 400
    user_id = request.args.get('user_id', 'default')
    try:
        luna = LunaAI()
        result = luna.generate_response(query, user_id)
        return jsonify(result)
    except Exception as e:
        return jsonify({'error': str(e)}), 500

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=7860, debug=True)