import os
import re
from pathlib import Path
from typing import List, Optional

import tmdbsimple as tmdb
from dotenv import load_dotenv
from models import (
    Collection,
    Genre,
    Movie,
    ProductionCompany,
    ProductionCountry,
    SpokenLanguage,
)
from sqlalchemy import extract
from sqlalchemy.exc import NoResultFound
from sqlmodel import Session, SQLModel, select

load_dotenv()

YEAR_PATTERN = re.compile(r"(\s\(\d{4}\))")


def create_model_obj(o: dict, model_type: SQLModel, session: Session) -> SQLModel:
    obj = model_type(**o)
    return obj


def create_model_objs(
    data: dict, model_type: SQLModel, session: Session
) -> List[SQLModel]:
    objs = []
    for o in data:
        obj = create_model_obj(o, model_type, session)
        objs.append(obj)
    return objs


def tmdb_info_to_movie(info: dict, session: Session) -> Movie:
    relationship_keys = {
        "genres",
        "belongs_to_collection",
        "production_companies",
        "production_countries",
        "spoken_languages",
    }
    movie_info = {k: v for k, v in info.items() if k not in relationship_keys}

    genres = create_model_objs(info["genres"], Genre, session)

    collection = None
    if info["belongs_to_collection"]:
        collection = create_model_obj(
            info["belongs_to_collection"], Collection, session
        )

    production_companies = create_model_objs(
        info["production_companies"], ProductionCompany, session
    )

    production_countries = create_model_objs(
        info["production_countries"], ProductionCountry, session
    )

    # languages
    spoken_languages = create_model_objs(
        info["spoken_languages"], SpokenLanguage, session
    )

    # create movie
    movie = Movie(**movie_info)
    movie.genres = genres
    movie.collection = collection
    movie.production_companies = production_companies
    movie.production_countries = production_countries
    movie.spoken_languages = spoken_languages

    session.add(movie)
    session.commit()
    session.refresh(movie)

    return movie


tmdb.API_KEY = os.getenv("TMDB_API_KEY", None)


def split_movie_path_title_and_year(path: str):
    movie_path = Path(path)
    movie_name = movie_path.stem

    year = None
    match = YEAR_PATTERN.search(movie_name)
    if match:
        year = match.group().strip(" ()")
        movie_name = movie_name.replace(match.group(), "")

    return movie_name, year


def get_movie_from_path(path: str, session: Session) -> Optional[Movie]:
    movie = None

    movie_name, year = split_movie_path_title_and_year(path)

    # lookup in db
    statement = select(Movie).where(Movie.title == movie_name)
    if year is not None:
        statement = statement.filter(extract("year", Movie.release_date) == int(year))
    try:
        movie = session.exec(statement).one()
    except NoResultFound:
        search = tmdb.Search()
        search.movie(query=movie_name, year=year)
        # take the first result:
        if search.results:
            id = search.results[0]["id"]
            info = tmdb.Movies(id).info()
            movie = tmdb_info_to_movie(info, session)

    return movie

re.compile dotenv.load_dotenv tmdbsimple.Search pathlib.Path tmdbsimple.Movies sqlalchemy.extract os.getenv models.Movie sqlmodel.select
[((384, 397), 'dotenv.load_dotenv', 'load_dotenv', ([], {}), '()\n', (395, 397), False, 'from dotenv import load_dotenv\n'), ((414, 445), 're.compile', 're.compile', (['"""(\\\\s\\\\(\\\\d{4}\\\\))"""'], {}), "('(\\\\s\\\\(\\\\d{4}\\\\))')\n", (424, 445), False, 'import re\n'), ((2087, 2118), 'os.getenv', 'os.getenv', (['"""TMDB_API_KEY"""', 'None'], {}), "('TMDB_API_KEY', None)\n", (2096, 2118), False, 'import os\n'), ((1746, 1765), 'models.Movie', 'Movie', ([], {}), '(**movie_info)\n', (1751, 1765), False, 'from models import Collection, Genre, Movie, ProductionCompany, ProductionCountry, SpokenLanguage\n'), ((2186, 2196), 'pathlib.Path', 'Path', (['path'], {}), '(path)\n', (2190, 2196), False, 'from pathlib import Path\n'), ((2625, 2638), 'sqlmodel.select', 'select', (['Movie'], {}), '(Movie)\n', (2631, 2638), False, 'from sqlmodel import Session, SQLModel, select\n'), ((2882, 2895), 'tmdbsimple.Search', 'tmdb.Search', ([], {}), '()\n', (2893, 2895), True, 'import tmdbsimple as tmdb\n'), ((2734, 2769), 'sqlalchemy.extract', 'extract', (['"""year"""', 'Movie.release_date'], {}), "('year', Movie.release_date)\n", (2741, 2769), False, 'from sqlalchemy import extract\n'), ((3066, 3081), 'tmdbsimple.Movies', 'tmdb.Movies', (['id'], {}), '(id)\n', (3077, 3081), True, 'import tmdbsimple as tmdb\n')]